diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 07329b6cd..6a516be2b 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -31,6 +31,7 @@ from __future__ import annotations import dataclasses import inspect +import json import logging from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple @@ -290,6 +291,42 @@ def attention_backends_of(cfg: Any) -> tuple: return prefill, decode +def modelexpress_transport_of(cfg: Any) -> str: + """The modelexpress transport a config-shaped object asks for. + + ``modelexpress_config`` is a JSON string (or an already-parsed dict) rather + than a leaf of its own; this is the shared parse for the transfer-engine + gate (`remote_instance_transfer_engine_of`) and any future bag reader. + ``ServerArgs.modelexpress_transport`` keeps its own instance-cached parse + (`_parsed_modelexpress_config`) -- same rule, cached seed-side.""" + raw = cfg.modelexpress_config + if raw is None: + parsed = {} + elif isinstance(raw, str): + parsed = json.loads(raw) + else: + parsed = raw + return parsed.get("transport", "nixl") + + +def remote_instance_transfer_engine_of(cfg: Any, load_format: Any = None) -> bool: + """Whether remote-instance weight loading runs over the transfer engine. + + ``load_format`` overrides the config's: a draft runner loading under + ``--speculative-draft-load-format`` needs its own transfer engine. Every + input is a ``model`` leaf, so this serves both the pre-publish member and + the post-publish accessor.""" + if cfg.remote_instance_weight_loader_start_seed_via_transfer_engine: + return True + if (load_format or cfg.load_format) != "remote_instance": + return False + backend = cfg.remote_instance_weight_loader_backend + return backend == "transfer_engine" or ( + backend == "modelexpress" + and modelexpress_transport_of(cfg) == "transfer_engine" + ) + + def mamba_extra_buffer_of(cfg: Any) -> bool: """Mid-resolution equivalent of runtime_context.mamba_extra_buffer_enabled: reads the (possibly overlaid) strategy from a config-shaped object. diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 09b0a9421..eb6041cd3 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1876,7 +1876,7 @@ class Scheduler( def process_input_requests(self, recv_reqs: List): now = time.monotonic() self.session_controller.maybe_reap(now) - if self.server_args.mm_feature_transport == "cuda_vmm": + if get_mm().mm_feature_transport == "cuda_vmm": for recv_req in recv_reqs: self._materialize_cuda_vmm_inputs(recv_req) diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 6e36069e5..1b4dfd78c 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -75,6 +75,7 @@ from sglang.srt.runtime_context import ( mamba_extra_buffer_enabled, mamba_extra_buffer_lazy_enabled, max_speculative_num_draft_tokens, + pre_capture_activation_reserve_mb, ) from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.spec_info import SpeculativeAlgorithm @@ -1782,7 +1783,7 @@ class KVCacheConfigurator: # Mamba state is a fixed pre-capture allocation, so it can't ride the ~0 post-capture slack. slack_gb = max( slack_gb, - self.server_args.pre_capture_activation_reserve_mb( + pre_capture_activation_reserve_mb( get_device_memory_capacity(self.device) ) / 1024, diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 7cf7c00ca..7f754bc24 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -179,6 +179,7 @@ from sglang.srt.runtime_context import ( get_spec, is_ep_joiner, is_ep_scale_joiner, + remote_instance_transfer_engine_enabled, set_global_dwdp_manager, ) from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo @@ -560,7 +561,6 @@ class ModelRunner: def init_remote_instance_weight_transporter(self): self.remote_instance_weight_transporter = RemoteInstanceWeightTransporter( - server_args=self.server_args, get_model=lambda: self.model, tp_rank=self.ps.tp_rank, gpu_id=self.gpu_id, @@ -671,9 +671,7 @@ class ModelRunner: ) def maybe_init_remote_instance_transfer_engine(self): - if self.server_args.remote_instance_weight_loader_use_transfer_engine( - load_format=self.draft_load_format - ): + if remote_instance_transfer_engine_enabled(load_format=self.draft_load_format): self.remote_instance_weight_transporter.init_engine() def maybe_init_expert_location_metadata(self): diff --git a/python/sglang/srt/model_executor/model_runner_components/kv_pool_runtime.py b/python/sglang/srt/model_executor/model_runner_components/kv_pool_runtime.py index 58c315149..9be8dfb43 100644 --- a/python/sglang/srt/model_executor/model_runner_components/kv_pool_runtime.py +++ b/python/sglang/srt/model_executor/model_runner_components/kv_pool_runtime.py @@ -11,6 +11,7 @@ from sglang.srt.distributed import get_world_group from sglang.srt.mem_cache.kv_cache_configurator import mm_runtime_reservation_gb from sglang.srt.model_executor.cuda_graph_config import Backend from sglang.srt.platforms import current_platform +from sglang.srt.runtime_context import pre_capture_activation_reserve_mb from sglang.srt.utils.common import get_available_gpu_memory, get_device_memory_capacity if TYPE_CHECKING: @@ -72,7 +73,7 @@ def compute_post_capture_kv_resize( if eager_decode_gap or mambaish_config(model_runner.model_config) is not None: headroom_gb = max( headroom_gb, - model_runner.server_args.pre_capture_activation_reserve_mb( + pre_capture_activation_reserve_mb( get_device_memory_capacity(model_runner.device) ) / 1024, diff --git a/python/sglang/srt/model_executor/model_runner_components/remote_instance_weight_transporter.py b/python/sglang/srt/model_executor/model_runner_components/remote_instance_weight_transporter.py index 5fbcab266..14690ab21 100644 --- a/python/sglang/srt/model_executor/model_runner_components/remote_instance_weight_transporter.py +++ b/python/sglang/srt/model_executor/model_runner_components/remote_instance_weight_transporter.py @@ -11,8 +11,11 @@ from sglang.srt.model_loader.remote_instance_weight_loader_utils import ( RemoteInstanceWeightLoaderBackend, register_memory_region, ) -from sglang.srt.runtime_context import get_model, get_parallel -from sglang.srt.server_args import ServerArgs +from sglang.srt.runtime_context import ( + get_model, + get_parallel, + remote_instance_transfer_engine_enabled, +) from sglang.srt.utils.network import NetworkAddress, get_local_ip_auto logger = logging.getLogger(__name__) @@ -20,7 +23,6 @@ logger = logging.getLogger(__name__) @dataclass(slots=True, kw_only=True) class RemoteInstanceWeightTransporter: - server_args: ServerArgs get_model: Callable[[], torch.nn.Module] tp_rank: int gpu_id: int @@ -55,7 +57,7 @@ class RemoteInstanceWeightTransporter: def maybe_register_and_publish_weight_info(self) -> None: if ( - self.server_args.remote_instance_weight_loader_use_transfer_engine() + remote_instance_transfer_engine_enabled() # ModelExpress owns TransferEngine memory registration and metadata # publishing for backend=modelexpress. Re-registering here would # overlap the same weight buffers. diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index b572581bd..a71ebff22 100644 --- a/python/sglang/srt/runtime_context.py +++ b/python/sglang/srt/runtime_context.py @@ -1420,6 +1420,53 @@ def mamba_extra_buffer_lazy_enabled() -> bool: ) +def remote_instance_transfer_engine_enabled(load_format: str | None = None) -> bool: + """Whether remote-instance weight loading runs over the transfer engine. + + Every input is a ``model`` leaf, so this derives from the bags and follows a + post-publish override; ``ServerArgs.remote_instance_weight_loader_use_transfer_engine`` + is the pre-publish equivalent, and both go through the same helper. + ``load_format`` is the caller's own (a draft runner loading under + ``--speculative-draft-load-format`` has one the process record does not). + """ + from sglang.srt.arg_groups.overrides import remote_instance_transfer_engine_of + + return remote_instance_transfer_engine_of(get_model(), load_format) + + +def pre_capture_activation_reserve_mb(gpu_mem: float | None) -> float: + """The activation working-set reserve held back before cuda-graph capture. + + Derived from published leaves across four bags (``disagg`` / ``schedule`` / + ``exec.graph`` / ``spec``) plus the configured parallel sizes, so it follows + a post-publish override; ``ServerArgs.pre_capture_activation_reserve_mb`` is + the pre-publish equivalent and + ``TestDerivedPredicatesAgreeAcrossTiers`` pins the two equal. + """ + schedule = get_schedule() + if get_disagg().disaggregation_mode == "decode": + running_requests = ( + schedule.max_running_requests + or get_exec().graph.cuda_graph_config.decode.max_bs + or 1 + ) + activation_tokens = max( + running_requests * (get_spec().speculative_num_draft_tokens or 1), 2048 + ) + elif schedule.chunked_prefill_size > 0: + activation_tokens = max(schedule.chunked_prefill_size, 2048) + else: + activation_tokens = max(schedule.max_prefill_tokens, 2048) + reserved_mem = ( + 512 + + activation_tokens * 1.5 + + _configured_parallel("tp_size") * _configured_parallel("pp_size") / 8 * 1024 + ) + if gpu_mem is not None and gpu_mem > 60 * 1024: + reserved_mem = max(reserved_mem, 10 * 1024) + return reserved_mem + + # --- Derived config accessors ------------------------------------------------ # # A few values are computed from several config fields plus the HF config, so diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 2c2813eb7..2def4ccb2 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -44,6 +44,7 @@ from sglang.srt.arg_groups.overrides import ( attention_backends_of, mamba_extra_buffer_lazy_of, mamba_extra_buffer_of, + remote_instance_transfer_engine_of, resolved_view, ) from sglang.srt.configs.embedding_model_spec import BCGPrefillPolicy @@ -9470,20 +9471,7 @@ class ServerArgs: def remote_instance_weight_loader_use_transfer_engine(self, load_format=None): """``load_format`` overrides the seed's: a draft runner loading under ``--speculative-draft-load-format`` needs its own transfer engine.""" - # Use TransferEngine as seed backend. - if self.remote_instance_weight_loader_start_seed_via_transfer_engine: - return True - # Use TransferEngine as client backend. - if (load_format or self.load_format) == "remote_instance" and ( - self.remote_instance_weight_loader_backend == "transfer_engine" - or ( - self.remote_instance_weight_loader_backend == "modelexpress" - and self.modelexpress_transport == "transfer_engine" - ) - ): - return True - else: - return False + return remote_instance_transfer_engine_of(self, load_format) def describe_kv_events_publisher(self) -> Optional[dict]: """Return a structured description of this server's KV-event diff --git a/python/sglang/srt/speculative/base_spec_worker.py b/python/sglang/srt/speculative/base_spec_worker.py index 259f9847f..6c2501515 100644 --- a/python/sglang/srt/speculative/base_spec_worker.py +++ b/python/sglang/srt/speculative/base_spec_worker.py @@ -11,7 +11,7 @@ from sglang.srt.model_executor.graph_memory_usage import ( merge_graph_memory_usage, merge_graph_time_usage, ) -from sglang.srt.runtime_context import get_exec, get_schedule +from sglang.srt.runtime_context import get_exec, get_memory, get_schedule if TYPE_CHECKING: from sglang.srt.managers.io_struct import ( @@ -235,7 +235,7 @@ class BaseSpecWorker(ABC): target_model_runner = self.target_worker.model_runner target_model_runner.mtp_draft_device_pools = () spec_algorithm = target_model_runner.spec_algorithm - if not self.server_args.enable_hierarchical_cache: + if not get_memory().enable_hierarchical_cache: return HiCacheDraftPlan() draft_runners = self._draft_model_runners() diff --git a/python/sglang/srt/speculative/draft_utils.py b/python/sglang/srt/speculative/draft_utils.py index bf3e92f29..821d28874 100644 --- a/python/sglang/srt/speculative/draft_utils.py +++ b/python/sglang/srt/speculative/draft_utils.py @@ -1,5 +1,4 @@ -from sglang.srt.runtime_context import get_spec -from sglang.srt.server_args import ServerArgs +from sglang.srt.runtime_context import attention_backends, get_spec from sglang.srt.utils.common import ( cpu_has_amx_support, is_blackwell, @@ -28,13 +27,11 @@ def _assert_draft_needs_no_conv_sidecar(draft_model_runner) -> None: class DraftBackendFactory: def __init__( self, - server_args: ServerArgs, draft_model_runner, topk: int, speculative_num_steps: int, seed_dsa_topk_from_draft_extend: bool = False, ): - self.server_args = server_args self.draft_model_runner = draft_model_runner self.topk = topk self.speculative_num_steps = speculative_num_steps @@ -45,13 +42,14 @@ class DraftBackendFactory: def _create_backend( self, backend_name: str, backend_map: dict, error_template: str ): - backend_type = ( - self.draft_attn_backend - if self.draft_attn_backend - else getattr(self.server_args, backend_name) + # The split pair with the base-backend fallback already applied. + prefill_backend, decode_backend = attention_backends() + configured = ( + decode_backend + if backend_name == "decode_attention_backend" + else prefill_backend ) - if backend_type is None: - backend_type = self.server_args.attention_backend + backend_type = self.draft_attn_backend or configured if backend_type not in backend_map: raise ValueError(error_template.format(backend_type=backend_type)) diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 2c2fc649e..9fb14f970 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -320,7 +320,6 @@ class EagleDraftWorker(EagleDraftWorkerBase): self.draft_extend_attn_backend = None draft_backend_factory = DraftBackendFactory( - self.server_args, self.draft_runner, self.topk, self.speculative_num_steps, diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index b9280f8a2..bf449b9fd 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -362,7 +362,6 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): self.draft_extend_attn_backend_list = [] for step in range(self.speculative_num_steps): draft_backend_factory = DraftBackendFactory( - self.server_args, self.draft_runner_list[step], self.topk, self.speculative_num_steps, diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py index b3e7dd235..e0df0d4b2 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py +++ b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py @@ -498,7 +498,6 @@ def _build_eagle_draft_extend_fixture( speculative_attention_mode="prefill", ) draft_extend_attn_backend = DraftBackendFactory( - fixture.runner.server_args, fixture.runner, settings.topk, settings.speculative_num_steps, diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py index 964eec4c8..dd76e9187 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py +++ b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py @@ -359,7 +359,6 @@ def _build_eagle_draft_fixture( ) _configure_runner_for_eagle_draft(fixture.runner, case, settings) draft_attn_backend = DraftBackendFactory( - fixture.runner.server_args, fixture.runner, settings.topk, settings.speculative_num_steps, diff --git a/test/registered/unit/multimodal/test_gpu_feature_transport.py b/test/registered/unit/multimodal/test_gpu_feature_transport.py index 287869e61..c2e7e57c4 100644 --- a/test/registered/unit/multimodal/test_gpu_feature_transport.py +++ b/test/registered/unit/multimodal/test_gpu_feature_transport.py @@ -490,6 +490,13 @@ class TestCudaVmmFeatureTransport(unittest.TestCase): class TestSchedulerMmTransportBoundary(unittest.TestCase): + def _publish(self, **fields): + from sglang.srt.runtime_context import get_context + + override = get_context().override_server_args(**fields) + override.install() + self.addCleanup(override.restore) + @staticmethod def _prepare_scheduler(scheduler): scheduler.session_controller = SimpleNamespace(maybe_reap=MagicMock()) @@ -501,7 +508,9 @@ class TestSchedulerMmTransportBoundary(unittest.TestCase): from sglang.srt.managers import scheduler as scheduler_module scheduler = object.__new__(scheduler_module.Scheduler) - scheduler.server_args = SimpleNamespace( + # The transport gate reads the published bags, so the case publishes + # the configuration under test. + self._publish( mm_feature_transport="cuda_vmm", enable_broadcast_mm_inputs_process=True, ) @@ -548,7 +557,7 @@ class TestSchedulerMmTransportBoundary(unittest.TestCase): return iter(self.batch) scheduler = object.__new__(scheduler_module.Scheduler) - scheduler.server_args = SimpleNamespace(mm_feature_transport="cuda_vmm") + self._publish(mm_feature_transport="cuda_vmm") self._prepare_scheduler(scheduler) raw_inputs = [object(), object()] materialized = [object(), object()] diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index 3d519b5cc..86e0f90af 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/test/registered/unit/test_runtime_context.py @@ -391,6 +391,19 @@ class _FakeResolvedArgs: speculative_num_draft_tokens: A[int | None, Arg(help="d"), NS("spec")] = None speculative_adaptive: A[bool, Arg(help="a"), NS("spec")] = False speculative_adaptive_config: A[str | None, Arg(help="c"), NS("spec")] = None + load_format: A[str, Arg(help="lf"), NS("model")] = "auto" + remote_instance_weight_loader_backend: A[str, Arg(help="rb"), NS("model")] = "nccl" + remote_instance_weight_loader_start_seed_via_transfer_engine: A[ + bool, Arg(help="rs"), NS("model") + ] = False + modelexpress_config: A[str | None, Arg(help="mx"), NS("model")] = None + disaggregation_mode: A[str, Arg(help="dm"), NS("disagg")] = "null" + max_running_requests: A[int | None, Arg(help="mrr"), NS("schedule")] = None + chunked_prefill_size: A[int, Arg(help="cps"), NS("schedule")] = -1 + max_prefill_tokens: A[int, Arg(help="mpt"), NS("schedule")] = 16384 + cuda_graph_config: A[object | None, Arg(help="cgc"), NS("exec.graph")] = None + tp_size: A[int, Arg(help="tp"), NS("parallel")] = 1 + pp_size: A[int, Arg(help="pp"), NS("parallel")] = 1 _resolved_overrides: list = dataclasses.field(default_factory=list) @@ -1013,6 +1026,74 @@ class TestDerivedPredicatesAgreeAcrossTiers(_IsolatedServerArgs): mamba_extra_buffer_lazy_enabled(), ) + def test_activation_reserve_matches_the_member(self): + from types import SimpleNamespace + + from sglang.srt.runtime_context import pre_capture_activation_reserve_mb + + graph = SimpleNamespace(decode=SimpleNamespace(max_bs=64)) + cases = ( + dict(disaggregation_mode="null", chunked_prefill_size=8192), + dict(disaggregation_mode="null", chunked_prefill_size=-1), + dict( + disaggregation_mode="null", + chunked_prefill_size=-1, + max_prefill_tokens=1024, + ), + dict(disaggregation_mode="decode", max_running_requests=32), + dict(disaggregation_mode="decode", max_running_requests=None), + dict( + disaggregation_mode="decode", + max_running_requests=None, + speculative_num_draft_tokens=4, + ), + dict( + disaggregation_mode="null", + chunked_prefill_size=8192, + tp_size=8, + pp_size=2, + ), + ) + for case in cases: + for gpu_mem in (None, 20 * 1024, 80 * 1024): + with self.subTest(gpu_mem=gpu_mem, **case): + args = _FakeResolvedArgs(cuda_graph_config=graph, **case) + get_context().set_server_args(args) + self.assertEqual( + ServerArgs.pre_capture_activation_reserve_mb(args, gpu_mem), + pre_capture_activation_reserve_mb(gpu_mem), + ) + + def test_remote_instance_transfer_engine_matches_the_member(self): + from sglang.srt.runtime_context import remote_instance_transfer_engine_enabled + + backends = ("nccl", "transfer_engine", "modelexpress") + transports = (None, '{"transport": "transfer_engine"}', '{"transport": "nixl"}') + for seed_via_te in (False, True): + for load_format in ("auto", "remote_instance"): + for backend in backends: + for mx in transports: + with self.subTest( + seed=seed_via_te, + load_format=load_format, + backend=backend, + modelexpress=mx, + ): + args = _FakeResolvedArgs( + load_format=load_format, + remote_instance_weight_loader_backend=backend, + remote_instance_weight_loader_start_seed_via_transfer_engine=seed_via_te, + modelexpress_config=mx, + ) + get_context().set_server_args(args) + for override in (None, "remote_instance", "auto"): + self.assertEqual( + ServerArgs.remote_instance_weight_loader_use_transfer_engine( + args, override + ), + remote_instance_transfer_engine_enabled(override), + ) + def test_attention_backends_match_the_member(self): from sglang.srt.runtime_context import attention_backends