Bound CP prefill batching by estimated temp memory

CP shared-KV bs>1 batching was only bounded by request count, extend tokens, and cached tokens. That left temporary GPU buffers such as MLA/index materialization, remap metadata, logits windows, and transfer descriptors implicit, and raw extend-token limits could exceed the active chunked-prefill budget.\n\nThis adds an explicit max-buffer-size admission gate with a CPU-only stream-aware estimator, wires it through PrefillAdder/Scheduler, performs a startup CUDA smoke allocation when configured, and reports the estimate in the scheduler admission benchmark. When chunked prefill is active, the effective CP extend-token limit is capped by the current chunk budget so the CP path does not advertise unreachable batch capacity or lift max-prefill-tokens too far.\n\nConstraint: Admission estimation must stay CPU-only on the scheduler hot path; CUDA allocation is limited to startup smoke checking.\nConstraint: Single oversized requests must still be allowed to run alone to avoid scheduler deadlock.\nRejected: Rely only on --max-prefill-tokens | it does not reliably bound the first oversized request and does not model cache-hit/load-back pressure.\nRejected: Let CP extend limit exceed chunked-prefill size | it creates an unreachable effective capacity and misleading budget lift.\nConfidence: medium\nScope-risk: moderate\nDirective: If bs>1 L1 prefetch is enabled later, update CPSharedKVPrefillBufferEstimatorContext.bs_gt1_l1_prefetch_enabled and include the live prefetch dense buffers in overlap windows.\nTested: local py_compile for touched files\nTested: local PYTHONPATH=python pytest -q test/registered/unit/managers/test_cp_shared_kv_prefill_buffer_estimator.py (4 passed)\nTested: remote cjy-glm5-new targeted pytest for new server_args, PrefillAdder, estimator, and benchmark cases (10 passed)\nTested: remote cjy-glm5-new PYTHONPATH=python pytest -q test/registered/unit/managers/test_cp_shared_kv_prefill_buffer_estimator.py test/registered/unit/managers/test_prefill_adder.py test/registered/unit/managers/test_prefill_scheduler_admission_bench.py (29 passed before chunk cap, then test_prefill_adder.py 21 passed after chunk cap)\nNot-tested: full server_args suite because existing TestPrepareServerArgs tries to reach HuggingFace and fails under container DNS/network\nNot-tested: GLM5 ETE smoke with --cp-shared-kv-prefill-max-buffer-size

python/sglang/srt/managers/cp_shared_kv_prefill_buffer_estimator.py

@@ -0,0 +1,249 @@
+"""CPU-only admission estimates for CP shared-KV prefill batching.
+
+The scheduler uses this module before it builds a real batch.  Keep this file
+free of CUDA allocations and collectives except for the explicit startup smoke
+check helper.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import Optional, Sequence
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+_DEFAULT_DESCRIPTOR_BYTES = 64
+_DEFAULT_LOGITS_DTYPE_BYTES = 2
+_DEFAULT_FALLBACK_KV_CACHE_DIM = 1
+_DEFAULT_FALLBACK_INDEX_HEAD_DIM = 128
+_DEFAULT_FALLBACK_QUANT_BLOCK_SIZE = 128
+
+
+@dataclass(frozen=True)
+class CPSharedKVPrefillBufferEstimatorContext:
+    kvcache: object | None
+    model_config: object | None
+    tp_size: int
+    page_size: int
+    logprob_chunk_enabled: bool
+    logprob_chunk_size: int
+    bs_gt1_l1_prefetch_enabled: bool = False
+
+
+@dataclass(frozen=True)
+class CPSharedKVPrefillBufferEstimate:
+    total_peak_bytes: int
+    layer_forward_peak_bytes: int
+    logits_window_peak_bytes: int
+    load_back_window_peak_bytes: int
+    materialize_peak_bytes: int
+    prefetch_peak_bytes: int
+    logits_peak_bytes: int
+    remap_peak_bytes: int
+    transfer_descriptor_peak_bytes: int
+    backup_descriptor_peak_bytes: int
+    l1_load_back_bytes: int
+    l1_extend_bytes: int
+    host_backup_bytes: int
+
+
+def ceil_paged_tokens(tokens: int, page_size: int) -> int:
+    if tokens <= 0:
+        return 0
+    return -(-int(tokens) // page_size) * page_size
+
+
+def ceil_pages(tokens: int, page_size: int) -> int:
+    if tokens <= 0:
+        return 0
+    return -(-int(tokens) // page_size)
+
+
+def _dtype_size(dtype: object | None, default: int = 2) -> int:
+    if dtype is None:
+        return default
+    itemsize = getattr(dtype, "itemsize", None)
+    if itemsize is not None:
+        return int(itemsize)
+    try:
+        return int(torch.empty((), dtype=dtype, device="cpu").element_size())
+    except Exception:
+        return default
+
+
+def _kv_cache_dim(kvcache: object | None) -> int:
+    return int(
+        getattr(kvcache, "kv_cache_dim", None)
+        or getattr(kvcache, "element_dim", None)
+        or _DEFAULT_FALLBACK_KV_CACHE_DIM
+    )
+
+
+def _kv_dtype_bytes(kvcache: object | None) -> int:
+    return _dtype_size(getattr(kvcache, "store_dtype", None), default=2)
+
+
+def _index_page_bytes(kvcache: object | None, page_size: int) -> int:
+    index_head_dim = int(
+        getattr(kvcache, "index_head_dim", None) or _DEFAULT_FALLBACK_INDEX_HEAD_DIM
+    )
+    quant_block_size = int(
+        getattr(kvcache, "quant_block_size", None)
+        or _DEFAULT_FALLBACK_QUANT_BLOCK_SIZE
+    )
+    index_dtype_bytes = _dtype_size(
+        getattr(kvcache, "index_k_with_scale_buffer_dtype", None),
+        default=1,
+    )
+    scale_bytes_per_token = (index_head_dim // quant_block_size) * 4
+    return page_size * (index_head_dim + scale_bytes_per_token) * index_dtype_bytes
+
+
+def _vocab_shard_size(model_config: object | None, tp_size: int) -> int:
+    vocab_size = int(getattr(model_config, "vocab_size", 0) or 0)
+    if vocab_size <= 0:
+        return 0
+    return -(-vocab_size // max(int(tp_size), 1))
+
+
+def estimate_cp_shared_kv_prefill_buffer_bytes(
+    *,
+    page_size: int,
+    batch_size: int,
+    prefix_lens: Sequence[int],
+    extend_lens: Sequence[int],
+    context: CPSharedKVPrefillBufferEstimatorContext | None,
+    return_logprob_rows: int = 0,
+    logits_dtype_bytes: int = _DEFAULT_LOGITS_DTYPE_BYTES,
+    descriptor_bytes: int = _DEFAULT_DESCRIPTOR_BYTES,
+) -> CPSharedKVPrefillBufferEstimate:
+    """Estimate peak temporary GPU buffer pressure for one candidate batch.
+
+    The estimate is intentionally conservative and stream-aware.  It sums
+    buffers that can be live concurrently on different streams instead of
+    taking a max over independent resource categories.
+    """
+
+    if context is None:
+        context = CPSharedKVPrefillBufferEstimatorContext(
+            kvcache=None,
+            model_config=None,
+            tp_size=1,
+            page_size=page_size,
+            logprob_chunk_enabled=False,
+            logprob_chunk_size=2048,
+            bs_gt1_l1_prefetch_enabled=False,
+        )
+
+    if page_size <= 0:
+        raise ValueError(f"page_size must be positive, got {page_size}")
+    if len(prefix_lens) != len(extend_lens):
+        raise ValueError(
+            "prefix_lens and extend_lens must have the same length: "
+            f"{len(prefix_lens)} != {len(extend_lens)}"
+        )
+
+    prefix_pages = [ceil_pages(tokens, page_size) for tokens in prefix_lens]
+    extend_pages = [ceil_pages(tokens, page_size) for tokens in extend_lens]
+    total_prefix_pages = sum(prefix_pages)
+    total_extend_pages = sum(extend_pages)
+    total_materialize_pages = sum(p + e for p, e in zip(prefix_pages, extend_pages))
+
+    kv_page_bytes = page_size * _kv_cache_dim(context.kvcache) * _kv_dtype_bytes(
+        context.kvcache
+    )
+    index_page_bytes = _index_page_bytes(context.kvcache, page_size)
+
+    materialize_peak_bytes = total_materialize_pages * (
+        kv_page_bytes + index_page_bytes
+    )
+    remap_peak_bytes = max(total_materialize_pages, int(batch_size)) * 8
+    prefetch_peak_bytes = (
+        total_prefix_pages * (kv_page_bytes + index_page_bytes)
+        if context.bs_gt1_l1_prefetch_enabled
+        else 0
+    )
+
+    logits_rows = max(int(batch_size), 0)
+    if return_logprob_rows > 0:
+        logits_rows = int(return_logprob_rows)
+        if context.logprob_chunk_enabled:
+            logits_rows = min(logits_rows, int(context.logprob_chunk_size))
+    vocab_shard = _vocab_shard_size(context.model_config, context.tp_size)
+    logits_peak_bytes = logits_rows * vocab_shard * int(logits_dtype_bytes)
+
+    transfer_descriptor_peak_bytes = total_prefix_pages * int(descriptor_bytes)
+    backup_descriptor_peak_bytes = total_extend_pages * int(descriptor_bytes)
+
+    l1_load_back_bytes = total_prefix_pages * kv_page_bytes
+    l1_extend_bytes = total_extend_pages * kv_page_bytes
+    host_backup_bytes = total_extend_pages * kv_page_bytes
+
+    layer_forward_peak_bytes = (
+        materialize_peak_bytes
+        + remap_peak_bytes
+        + prefetch_peak_bytes
+        + backup_descriptor_peak_bytes
+    )
+    logits_window_peak_bytes = (
+        logits_peak_bytes + prefetch_peak_bytes + backup_descriptor_peak_bytes
+    )
+    load_back_window_peak_bytes = (
+        transfer_descriptor_peak_bytes
+        + prefetch_peak_bytes
+        + backup_descriptor_peak_bytes
+    )
+    total_peak_bytes = max(
+        layer_forward_peak_bytes,
+        logits_window_peak_bytes,
+        load_back_window_peak_bytes,
+    )
+
+    return CPSharedKVPrefillBufferEstimate(
+        total_peak_bytes=total_peak_bytes,
+        layer_forward_peak_bytes=layer_forward_peak_bytes,
+        logits_window_peak_bytes=logits_window_peak_bytes,
+        load_back_window_peak_bytes=load_back_window_peak_bytes,
+        materialize_peak_bytes=materialize_peak_bytes,
+        prefetch_peak_bytes=prefetch_peak_bytes,
+        logits_peak_bytes=logits_peak_bytes,
+        remap_peak_bytes=remap_peak_bytes,
+        transfer_descriptor_peak_bytes=transfer_descriptor_peak_bytes,
+        backup_descriptor_peak_bytes=backup_descriptor_peak_bytes,
+        l1_load_back_bytes=l1_load_back_bytes,
+        l1_extend_bytes=l1_extend_bytes,
+        host_backup_bytes=host_backup_bytes,
+    )
+
+
+def smoke_check_cp_shared_kv_prefill_buffer_size(
+    *,
+    device: torch.device | str,
+    size_bytes: int,
+    torch_module=torch,
+) -> None:
+    """Fail fast if startup cannot allocate the configured temp buffer size."""
+
+    probe = None
+    try:
+        torch_module.cuda.synchronize()
+        probe = torch_module.empty(size_bytes, dtype=torch_module.uint8, device=device)
+        if size_bytes > 0:
+            probe[0] = 0
+            probe[-1] = 0
+        torch_module.cuda.synchronize()
+    except torch_module.cuda.OutOfMemoryError as exc:
+        raise RuntimeError(
+            "[CP_SHARED_KV_FAIL_FAST][prefill_buffer_smoke] "
+            f"cannot allocate cp_shared_kv_prefill_max_buffer_size={size_bytes} "
+            "bytes after scheduler startup. Lower "
+            "--cp-shared-kv-prefill-max-buffer-size or reduce other static GPU "
+            "memory consumers."
+        ) from exc
+    finally:
+        del probe
+        torch_module.cuda.empty_cache()

python/sglang/srt/managers/schedule_policy.py

@@ -34,8 +34,14 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Set, Union
 import torch
 
 from sglang.srt.dllm.config import DllmConfig
+from sglang.srt.environ import envs
 from sglang.srt.layers.attention.nsa.utils import is_nsa_prefill_cp_in_seq_split
 from sglang.srt.layers.utils.cp_utils import is_prefill_context_parallel_enabled
+from sglang.srt.managers.cp_shared_kv_prefill_buffer_estimator import (
+    CPSharedKVPrefillBufferEstimate,
+    CPSharedKVPrefillBufferEstimatorContext,
+    estimate_cp_shared_kv_prefill_buffer_bytes,
+)
 from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
 from sglang.srt.mem_cache.base_prefix_cache import (
     BasePrefixCache,
@@ -77,6 +83,24 @@ IN_BATCH_PREFIX_CACHING_DEPRIORITIZE_THRESHOLD = int(
 IGNORE_EOS_RESERVE_TOKENS = 1
 
 
+_CP_SHARED_KV_BS_GT1_ADMISSION_DEBUG_COUNTS: Dict[str, int] = {}
+
+
+def _cp_shared_kv_bs_gt1_admission_debug(
+    key: str,
+    message: str,
+    *args,
+) -> None:
+    if not envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.get():
+        return
+    limit = int(envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG_LIMIT.get())
+    count = _CP_SHARED_KV_BS_GT1_ADMISSION_DEBUG_COUNTS.get(key, 0)
+    if limit > 0 and count >= limit:
+        return
+    _CP_SHARED_KV_BS_GT1_ADMISSION_DEBUG_COUNTS[key] = count + 1
+    logger.info("[CP_SHARED_KV_BS_GT1_DEBUG] event=%s " + message, key, *args)
+
+
 class CacheAwarePolicy(Enum):
     """Scheduling policies that are aware of the tree cache."""
 
@@ -394,6 +418,10 @@ class PrefillAdder:
         cp_shared_kv_prefill_max_batch_requests: Optional[int] = None,
         cp_shared_kv_prefill_max_total_extend_tokens: Optional[int] = None,
         cp_shared_kv_prefill_max_total_cached_tokens: Optional[int] = None,
+        cp_shared_kv_prefill_max_buffer_size: Optional[int] = None,
+        cp_shared_kv_prefill_buffer_estimator_context: Optional[
+            CPSharedKVPrefillBufferEstimatorContext
+        ] = None,
         prefill_delayer_single_pass: Optional[PrefillDelayerSinglePassExecutor] = None,
         dllm_config: Optional[DllmConfig] = None,
     ):
@@ -452,9 +480,30 @@ class PrefillAdder:
         self.cp_shared_kv_prefill_max_total_extend_tokens = (
             cp_shared_kv_prefill_max_total_extend_tokens
         )
+        if (
+            self._is_cp_prefill_context()
+            and self.enable_cp_shared_kv_prefill_bs_gt1
+            and self.cp_shared_kv_prefill_max_total_extend_tokens is not None
+            and self.rem_chunk_tokens is not None
+        ):
+            # When chunked prefill is active, a batch cannot consume more
+            # extend tokens than the current chunk budget.  Use the smaller
+            # value as the effective CP-specific grouping limit so the generic
+            # max_prefill_tokens lift below does not advertise unreachable
+            # capacity.
+            self.cp_shared_kv_prefill_max_total_extend_tokens = min(
+                self.cp_shared_kv_prefill_max_total_extend_tokens,
+                max(self.rem_chunk_tokens, 0),
+            )
         self.cp_shared_kv_prefill_max_total_cached_tokens = (
             cp_shared_kv_prefill_max_total_cached_tokens
         )
+        self.cp_shared_kv_prefill_max_buffer_size = (
+            cp_shared_kv_prefill_max_buffer_size
+        )
+        self.cp_shared_kv_prefill_buffer_estimator_context = (
+            cp_shared_kv_prefill_buffer_estimator_context
+        )
         if (
             self._is_cp_prefill_context()
             and self.enable_cp_shared_kv_prefill_bs_gt1
@@ -479,6 +528,12 @@ class PrefillAdder:
             )
         self.cp_shared_kv_prefill_total_extend_tokens = 0
         self.cp_shared_kv_prefill_total_cached_tokens = 0
+        self.cp_shared_kv_prefill_estimate_prefix_lens: list[int] = []
+        self.cp_shared_kv_prefill_estimate_extend_lens: list[int] = []
+        self.cp_shared_kv_prefill_estimated_peak_buffer_bytes = 0
+        self.cp_shared_kv_prefill_last_buffer_estimate: Optional[
+            CPSharedKVPrefillBufferEstimate
+        ] = None
 
     def _init_dllm_meta(self, dllm_config: DllmConfig):
         self.dllm_block_size = dllm_config.block_size
@@ -599,6 +654,68 @@ class PrefillAdder:
         # grouping pressure from prefix/load-back work, not legal request size.
         return projected > limit and len(self.can_run_list) > 0
 
+    def _estimate_cp_prefill_projected_buffer(
+        self, prefix_len: int, extend_input_len: int
+    ) -> CPSharedKVPrefillBufferEstimate:
+        return estimate_cp_shared_kv_prefill_buffer_bytes(
+            page_size=self.page_size,
+            batch_size=len(self.cp_shared_kv_prefill_estimate_prefix_lens) + 1,
+            prefix_lens=[
+                *self.cp_shared_kv_prefill_estimate_prefix_lens,
+                prefix_len,
+            ],
+            extend_lens=[
+                *self.cp_shared_kv_prefill_estimate_extend_lens,
+                extend_input_len,
+            ],
+            context=self.cp_shared_kv_prefill_buffer_estimator_context,
+        )
+
+    def _cp_prefill_buffer_limit_exceeded(
+        self, prefix_len: int, extend_input_len: int, rid: Optional[str] = None
+    ) -> bool:
+        if not (
+            self._is_cp_prefill_context()
+            and self.enable_cp_shared_kv_prefill_bs_gt1
+        ):
+            return False
+
+        limit = self.cp_shared_kv_prefill_max_buffer_size
+        if limit is None:
+            return False
+
+        estimate = self._estimate_cp_prefill_projected_buffer(
+            prefix_len, extend_input_len
+        )
+        self.cp_shared_kv_prefill_last_buffer_estimate = estimate
+
+        # Do not deadlock a single large request. The limit bounds grouping, not
+        # the maximum legal request size; actual runtime capacity remains
+        # allocator/CUDA-owned.
+        should_stop = estimate.total_peak_bytes > limit and len(self.can_run_list) > 0
+        if should_stop:
+            _cp_shared_kv_bs_gt1_admission_debug(
+                "admission_max_buffer_size",
+                "stop rid=%s projected=%s limit=%s batch_size=%s "
+                "layer_forward=%s logits_window=%s load_back_window=%s "
+                "materialize=%s prefetch=%s logits=%s remap=%s "
+                "transfer_desc=%s backup_desc=%s",
+                rid,
+                estimate.total_peak_bytes,
+                limit,
+                len(self.can_run_list) + 1,
+                estimate.layer_forward_peak_bytes,
+                estimate.logits_window_peak_bytes,
+                estimate.load_back_window_peak_bytes,
+                estimate.materialize_peak_bytes,
+                estimate.prefetch_peak_bytes,
+                estimate.logits_peak_bytes,
+                estimate.remap_peak_bytes,
+                estimate.transfer_descriptor_peak_bytes,
+                estimate.backup_descriptor_peak_bytes,
+            )
+        return should_stop
+
     def _get_available_device_tokens_for_load_back(self) -> int:
         if self.is_hybrid_swa:
             return (
@@ -648,6 +765,20 @@ class PrefillAdder:
         self.rem_input_tokens -= extend_input_len
         if self._is_cp_prefill_context():
             self.cp_shared_kv_prefill_total_extend_tokens += extend_input_len
+            if self.cp_shared_kv_prefill_max_buffer_size is not None:
+                self.cp_shared_kv_prefill_estimate_prefix_lens.append(prefix_len)
+                self.cp_shared_kv_prefill_estimate_extend_lens.append(extend_input_len)
+                estimate = estimate_cp_shared_kv_prefill_buffer_bytes(
+                    page_size=self.page_size,
+                    batch_size=len(self.cp_shared_kv_prefill_estimate_prefix_lens),
+                    prefix_lens=self.cp_shared_kv_prefill_estimate_prefix_lens,
+                    extend_lens=self.cp_shared_kv_prefill_estimate_extend_lens,
+                    context=self.cp_shared_kv_prefill_buffer_estimator_context,
+                )
+                self.cp_shared_kv_prefill_last_buffer_estimate = estimate
+                self.cp_shared_kv_prefill_estimated_peak_buffer_bytes = (
+                    estimate.total_peak_bytes
+                )
 
         if self.dllm_config is not None:
             self.rem_dllm_tokens -= extend_input_len
@@ -934,6 +1065,10 @@ class PrefillAdder:
                 return AddReqResult.OTHER
             if self._cp_prefill_cached_limit_exceeded(prefix_len):
                 return AddReqResult.OTHER
+            if self._cp_prefill_buffer_limit_exceeded(
+                prefix_len, input_tokens, req.rid
+            ):
+                return AddReqResult.OTHER
 
             if self.dllm_config is not None:
                 if self.rem_dllm_tokens <= 0:
@@ -980,6 +1115,10 @@ class PrefillAdder:
                     return AddReqResult.OTHER
                 if self._cp_prefill_cached_limit_exceeded(prefix_len):
                     return AddReqResult.OTHER
+                if self._cp_prefill_buffer_limit_exceeded(
+                    prefix_len, trunc_len, req.rid
+                ):
+                    return AddReqResult.OTHER
 
                 # Chunked prefill
                 req.set_extend_input_len(trunc_len)

python/sglang/srt/managers/scheduler.py

@@ -147,6 +147,10 @@ from sglang.srt.managers.prefill_delayer import (
     PrefillDelayer,
     PrefillDelayerSinglePassExecutor,
 )
+from sglang.srt.managers.cp_shared_kv_prefill_buffer_estimator import (
+    CPSharedKVPrefillBufferEstimatorContext,
+    smoke_check_cp_shared_kv_prefill_buffer_size,
+)
 from sglang.srt.managers.schedule_batch import (
     FINISH_ABORT,
     ModelWorkerBatch,
@@ -492,6 +496,8 @@ class Scheduler(
         # Init the grammar backend for constrained generation
         self.grammar_manager = GrammarManager(self)
 
+        self.maybe_smoke_check_cp_shared_kv_prefill_buffer()
+
         self.is_initializing = False
 
     def init_model_config(self):
@@ -515,6 +521,62 @@ class Scheduler(
                     )
                     self.page_size = self.dllm_config.block_size
 
+    def make_cp_shared_kv_prefill_buffer_estimator_context(
+        self,
+    ) -> CPSharedKVPrefillBufferEstimatorContext:
+        return CPSharedKVPrefillBufferEstimatorContext(
+            kvcache=self.token_to_kv_pool_allocator.get_kvcache(),
+            model_config=self.model_config,
+            tp_size=self.tp_size,
+            page_size=self.page_size,
+            logprob_chunk_enabled=envs.SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK.get(),
+            logprob_chunk_size=envs.SGLANG_LOGITS_PROCESSER_CHUNK_SIZE.get(),
+            # As of this scheduler path, CP shared-KV L1 prefetch explicitly
+            # gates off bs>1. Keep the estimate aligned with runtime until that
+            # path is enabled.
+            bs_gt1_l1_prefetch_enabled=False,
+        )
+
+    def maybe_smoke_check_cp_shared_kv_prefill_buffer(self) -> None:
+        size_bytes = self.server_args.cp_shared_kv_prefill_max_buffer_size
+        if not (
+            self.server_args.enable_cp_shared_kv_prefill_bs_gt1
+            and self.server_args.enable_nsa_prefill_cp_shared_kv
+            and size_bytes is not None
+            and self.server_args.disaggregation_mode in (None, "null", "prefill")
+        ):
+            return
+
+        if not str(self.device).startswith("cuda"):
+            logger.warning(
+                "[CP_SHARED_KV_FALLBACK][prefill_buffer_smoke] "
+                "skip CUDA smoke allocation on non-CUDA device=%s size_bytes=%s",
+                self.device,
+                size_bytes,
+            )
+            return
+
+        logger.info(
+            "[CP_SHARED_KV_BS_GT1] prefill buffer smoke allocation begin: "
+            "rank=%s cp_rank=%s size_bytes=%s size_gb=%.3f",
+            self.tp_rank,
+            self.attn_cp_rank,
+            size_bytes,
+            size_bytes / 1e9,
+        )
+        smoke_check_cp_shared_kv_prefill_buffer_size(
+            device=self.device,
+            size_bytes=size_bytes,
+        )
+        logger.info(
+            "[CP_SHARED_KV_BS_GT1] prefill buffer smoke allocation passed: "
+            "rank=%s cp_rank=%s size_bytes=%s size_gb=%.3f",
+            self.tp_rank,
+            self.attn_cp_rank,
+            size_bytes,
+            size_bytes / 1e9,
+        )
+
     def init_ipc_channels(self, port_args: PortArgs):
         context = zmq.Context(2)
         self.idle_sleeper = None
@@ -2422,6 +2484,14 @@ class Scheduler(
             cp_shared_kv_prefill_max_total_cached_tokens=(
                 self.server_args.cp_shared_kv_prefill_max_total_cached_tokens
             ),
+            cp_shared_kv_prefill_max_buffer_size=(
+                self.server_args.cp_shared_kv_prefill_max_buffer_size
+            ),
+            cp_shared_kv_prefill_buffer_estimator_context=(
+                self.make_cp_shared_kv_prefill_buffer_estimator_context()
+                if self.server_args.cp_shared_kv_prefill_max_buffer_size is not None
+                else None
+            ),
             prefill_delayer_single_pass=prefill_delayer_single_pass,
             dllm_config=self.dllm_config,
         )

python/sglang/srt/server_args.py

@@ -23,7 +23,9 @@ import json
 import logging
 import os
 import random
+import re
 import tempfile
+from decimal import Decimal
 from typing import Any, Callable, Dict, List, Literal, Optional, Union
 
 from sglang.srt.connector import ConnectorType
@@ -72,6 +74,23 @@ from sglang.utils import is_in_ci
 
 logger = logging.getLogger(__name__)
 
+
+def human_readable_gb_size(value: str) -> int:
+    """Parse sizes where a bare number means decimal GB.
+
+    Examples:
+        '8' -> 8000000000
+        '8G' -> 8000000000
+        '8Gi' -> 8589934592
+
+    Explicit suffixes reuse human_readable_int semantics.
+    """
+
+    value = value.strip()
+    if re.fullmatch(r"\d+(?:\.\d+)?", value):
+        return int(Decimal(value) * Decimal(10**9))
+    return human_readable_int(value)
+
 # Define constants
 DEFAULT_UVICORN_ACCESS_LOG_EXCLUDE_PREFIXES = ()
 SAMPLING_BACKEND_CHOICES = {"flashinfer", "pytorch", "ascend"}
@@ -678,6 +697,7 @@ class ServerArgs:
     cp_shared_kv_prefill_max_batch_requests: Optional[int] = None
     cp_shared_kv_prefill_max_total_extend_tokens: Optional[int] = None
     cp_shared_kv_prefill_max_total_cached_tokens: Optional[int] = None
+    cp_shared_kv_prefill_max_buffer_size: Optional[int] = None
     enable_fused_qk_norm_rope: bool = False
     enable_precise_embedding_interpolation: bool = False
     enable_fused_moe_sum_all_reduce: bool = False
@@ -1002,6 +1022,14 @@ class ServerArgs:
                 "cp_shared_kv_prefill_max_total_cached_tokens must be a positive "
                 "integer when specified."
             )
+        if (
+            self.cp_shared_kv_prefill_max_buffer_size is not None
+            and self.cp_shared_kv_prefill_max_buffer_size <= 0
+        ):
+            raise ValueError(
+                "cp_shared_kv_prefill_max_buffer_size must be a positive "
+                "integer when specified."
+            )
 
     def _handle_cp_hicache_layout_validation(self):
         if not (
@@ -5919,6 +5947,20 @@ class ServerArgs:
             )
             + f"\n\n{human_readable_int.__doc__}",
         )
+        parser.add_argument(
+            "--cp-shared-kv-prefill-max-buffer-size",
+            type=human_readable_gb_size,
+            default=ServerArgs.cp_shared_kv_prefill_max_buffer_size,
+            help=(
+                "Maximum estimated peak temporary GPU buffer size admitted into "
+                "one NSA in-seq CP shared-KV prefill batch when "
+                "--enable-cp-shared-kv-prefill-bs-gt1 is set. Bare numbers are "
+                "interpreted as decimal GB (for example, 8 means 8G). Explicit "
+                "SI/IEC suffixes such as 8G or 8Gi are also accepted. A single "
+                "request larger than this limit is still allowed to run alone."
+            )
+            + f"\n\n{human_readable_gb_size.__doc__}",
+        )
         parser.add_argument(
             "--nsa-prefill-cp-mode",
             type=str,