Add a provider hook for prefill-buffer ceilings (#39182)

Co-authored-by: cctry <17473714+cctry@users.noreply.github.com>

python/sglang/srt/arg_groups/overrides.py

@@ -68,6 +68,7 @@ from sglang.srt.arg_groups.model_override_base import (  # noqa: F401
     resolving_view,
     use_mla_backend,
 )
+from sglang.srt.arg_groups.prefill_buffer_ceiling import prefill_buffer_ceiling_of
 
 logger = logging.getLogger(__name__)
 from sglang.srt.environ import envs
@@ -1873,7 +1874,10 @@ def cutedsl_moe_max_num_tokens(server_args: Any) -> int:
 
 def max_prefill_buffer_tokens(server_args: Any) -> int:
     """Prefill-buffer ceiling: chunked_prefill_size, except PP dynamic
-    chunking can grow chunks toward max_prefill_tokens and probe at 1.25x."""
+    chunking can grow chunks toward max_prefill_tokens and probe at 1.25x.
+
+    Records with a registered ceiling provider (see
+    ``register_prefill_buffer_ceiling``) answer through it."""
     cfg = resolving_view(server_args)
     chunked = (
         cfg.chunked_prefill_size
@@ -1883,7 +1887,10 @@ def max_prefill_buffer_tokens(server_args: Any) -> int:
     tokens = chunked
     if cfg.enable_dynamic_chunking and cfg.pp_size > 1 and chunked:
         tokens = max(tokens, cfg.max_prefill_tokens or 0, math.ceil(chunked * 1.25))
-    return tokens
+    record = server_args
+    if isinstance(server_args, (ResolvedView, ResolvingConfig)):
+        record = record_of(server_args)
+    return prefill_buffer_ceiling_of(record, tokens)
 
 
 def mamba_cache_chunk_size(server_args: Any) -> int:

python/sglang/srt/arg_groups/prefill_buffer_ceiling.py

@@ -0,0 +1,35 @@
+"""Dependency-free hook shared by pre- and post-publish prefill-buffer sizing."""
+
+from typing import Any, Callable, Optional
+
+_prefill_buffer_ceiling_fn: Optional[Callable[[Any, int], int]] = None
+
+
+def register_prefill_buffer_ceiling(
+    fn: Callable[[Any, int], int],
+) -> Callable[[Any, int], int]:
+    """Register one provider; repeating the same registration is harmless.
+
+    The provider receives ``(record, default_ceiling)`` and returns the ceiling,
+    keeping ``default_ceiling`` for records it does not handle. ``record`` is
+    the original argument record, never a view: its fields remain raw inputs.
+    Read resolution declarations through ``resolving_view(record)`` from
+    ``arg_groups.model_override_base``. The provider must not mutate the record
+    and must use the same sizing policy before and after publication.
+
+    Registering a different provider raises instead of replacing the first.
+    """
+    global _prefill_buffer_ceiling_fn
+    if _prefill_buffer_ceiling_fn is not None and _prefill_buffer_ceiling_fn is not fn:
+        raise RuntimeError(
+            "A different prefill-buffer ceiling provider is already registered"
+        )
+    _prefill_buffer_ceiling_fn = fn
+    return fn
+
+
+def prefill_buffer_ceiling_of(record: Any, default_ceiling: int) -> int:
+    """Return the provider's ceiling, or the default when none is registered."""
+    if _prefill_buffer_ceiling_fn is None:
+        return default_ceiling
+    return _prefill_buffer_ceiling_fn(record, default_ceiling)

python/sglang/srt/runtime_context.py

@@ -57,6 +57,8 @@ from typing import TYPE_CHECKING, Any, Dict, Optional
 
 import msgspec
 
+from sglang.srt.arg_groups.prefill_buffer_ceiling import prefill_buffer_ceiling_of
+
 if TYPE_CHECKING:
     from sglang.srt.model_executor.runner_utils.pool import GraphPoolBorrowState
     from sglang.srt.server_args import ServerArgs
@@ -1792,13 +1794,13 @@ def max_prefill_buffer_tokens() -> int:
     """The prefill-buffer ceiling: ``chunked_prefill_size``, except PP dynamic
     chunking can grow chunks toward ``max_prefill_tokens`` and probe at 1.25x.
 
-    Every input is a published leaf (``schedule`` plus the configured PP size),
-    so this derives from the bags and follows a post-publish override;
+    The default derives from published leaves (``schedule`` plus the configured
+    PP size), so it follows post-publish overrides;
     ``overrides.max_prefill_buffer_tokens`` is the pre-publish equivalent and
-    ``TestDerivedPredicatesAgreeAcrossTiers`` pins the two equal.
+    ``TestDerivedPredicatesAgreeAcrossTiers`` pins the two equal. Records with
+    a registered ceiling provider (see ``register_prefill_buffer_ceiling``)
+    answer through it.
     """
-    import math
-
     schedule = get_schedule()
     chunked = (
         schedule.chunked_prefill_size
@@ -1810,7 +1812,7 @@ def max_prefill_buffer_tokens() -> int:
         tokens = max(
             tokens, schedule.max_prefill_tokens or 0, math.ceil(chunked * 1.25)
         )
-    return tokens
+    return prefill_buffer_ceiling_of(get_server_args(), tokens)
 
 
 def pre_capture_activation_reserve_mb(gpu_mem: float | None) -> float: