diff --git a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py index b7d450753..600877617 100644 --- a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py +++ b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py @@ -21,6 +21,7 @@ from sglang.srt.mem_cache.memory_pool import ( from sglang.srt.mem_cache.pool_host.common import get_allocator_type from sglang.srt.mem_cache.pool_host.mha import get_mha_host_pool_cls from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost +from sglang.srt.runtime_context import get_schedule from sglang.srt.server_args import ServerArgs from sglang.srt.utils.common import ceil_align @@ -43,8 +44,7 @@ class DecodeKVCacheOffloadManager: ) -> None: self.req_to_token_pool = req_to_token_pool self.token_to_kv_pool_allocator = token_to_kv_pool_allocator - self.page_size = server_args.page_size - self.server_args = server_args + self.page_size = get_schedule().page_size self.request_counter = 0 self.tree_cache = tree_cache env_stride = envs.SGLANG_HICACHE_DECODE_OFFLOAD_STRIDE.get() diff --git a/python/sglang/srt/eplb/expert_distribution.py b/python/sglang/srt/eplb/expert_distribution.py index d23703f7c..e28f583ee 100644 --- a/python/sglang/srt/eplb/expert_distribution.py +++ b/python/sglang/srt/eplb/expert_distribution.py @@ -35,6 +35,7 @@ from sglang.srt.observability.metrics_collector import ( ExpertDispatchCollector, resolve_collector_class, ) +from sglang.srt.runtime_context import get_schedule from sglang.srt.server_args import ServerArgs from sglang.srt.utils import Withable, get_device, get_int_env_var @@ -390,7 +391,7 @@ class _DetailSinglePassGatherer(_SinglePassGatherer): ( expert_location_metadata.num_layers, # TODO determine the max number - server_args.chunked_prefill_size * 8, + get_schedule().chunked_prefill_size * 8, self._TOP_K_NUM, ), dtype=torch.int32, diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py index 5f06cdafc..815c70740 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py @@ -16,7 +16,7 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.environ import envs from sglang.srt.layers.deep_gemm_wrapper.configurer import ENABLE_JIT_DEEPGEMM from sglang.srt.model_executor.forward_batch_info import ForwardMode -from sglang.srt.runtime_context import get_parallel +from sglang.srt.runtime_context import get_parallel, get_schedule from sglang.srt.server_args import ServerArgs from sglang.srt.utils import ceil_align, ceil_div, get_available_gpu_memory, is_musa @@ -67,9 +67,10 @@ def update_deep_gemm_config(gpu_id: int, server_args: ServerArgs): # 8192, 9008, ... 16384 (step 16) # Totally 1024 + 1024 / 2 + 2048 / 4 + 4096 / 8 + 8192 / 16 = 3072 kernels next_m, sample_step = 1024, 2 + chunked_prefill_size = get_schedule().chunked_prefill_size max_prefill_bs = ( - min(server_args.chunked_prefill_size, 32 * 1024) - if server_args.chunked_prefill_size >= 1 + min(chunked_prefill_size, 32 * 1024) + if chunked_prefill_size >= 1 else 16 * 1024 ) while next_m < max_prefill_bs: @@ -81,10 +82,11 @@ def update_deep_gemm_config(gpu_id: int, server_args: ServerArgs): else: # When fast warmup isn't enabled, generate m_max and compile all the covered Ms. m_max = 1024 * 16 - if server_args.chunked_prefill_size < 1: + chunked_prefill_size = get_schedule().chunked_prefill_size + if chunked_prefill_size < 1: m_max = 1024 * 64 - elif server_args.chunked_prefill_size > 8192: - m_max = server_args.chunked_prefill_size * 2 + elif chunked_prefill_size > 8192: + m_max = chunked_prefill_size * 2 m_max = min(1024 * 128, m_max) _BUILTIN_M_LIST += list(range(1, m_max + 1)) diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py index 054a067ac..dc9316806 100644 --- a/python/sglang/srt/lora/lora_manager.py +++ b/python/sglang/srt/lora/lora_manager.py @@ -46,7 +46,7 @@ from sglang.srt.lora.utils import ( ) from sglang.srt.managers.io_struct import LoRAUpdateOutput from sglang.srt.model_executor.forward_batch_info import ForwardBatch -from sglang.srt.runtime_context import get_parallel +from sglang.srt.runtime_context import get_exec, get_parallel from sglang.srt.server_args import ServerArgs from sglang.srt.utils import get_available_gpu_memory, replace_submodule from sglang.srt.utils.hf_transformers_utils import AutoConfig @@ -1030,7 +1030,7 @@ def init_lora_cuda_graph_moe_buffers( """ from sglang.srt.lora.layers import FusedMoEWithLoRA - max_bs = server_args.cuda_graph_config.decode.max_bs + max_bs = get_exec().graph.cuda_graph_config.decode.max_bs max_loras = server_args.max_loras_per_batch for module in model.modules(): if isinstance(module, FusedMoEWithLoRA): diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 9a15cb975..f3e73e7a9 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -436,7 +436,7 @@ class Scheduler( self.spec_algorithm = SpeculativeAlgorithm.from_string( server_args.speculative_algorithm ) - self.page_size = server_args.page_size + self.page_size = get_schedule().page_size self.enable_hierarchical_cache = server_args.enable_hierarchical_cache self.enable_session_radix_cache = server_args.enable_session_radix_cache self.enable_hicache_storage = server_args.hicache_storage_backend is not None diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index cdcc2b5c7..a1459d982 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/python/sglang/srt/mem_cache/kv_cache_builder.py @@ -36,7 +36,7 @@ from sglang.srt.managers.mm_schedule import init_mm_embedding_cache from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.registry import TreeCacheBuildContext, create_tree_cache from sglang.srt.model_loader.utils import get_resolved_model_impl -from sglang.srt.runtime_context import get_parallel +from sglang.srt.runtime_context import get_parallel, get_schedule if TYPE_CHECKING: @@ -205,7 +205,7 @@ def build_kv_cache( "with Mamba/SSM models" ) - effective_chunked_prefill_size = server_args.chunked_prefill_size + effective_chunked_prefill_size = get_schedule().chunked_prefill_size if model_config.is_multimodal and uses_transformers_backend: effective_chunked_prefill_size = None diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 7f754bc24..72b9c8b3d 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -313,6 +313,13 @@ class ModelRunner: self.dist_port = nccl_port self.server_args = server_args self.is_draft_worker = is_draft_worker + # Set the global server_args in the scheduler process (target worker + # only, so a draft init cannot clobber target-derived global state). + # Before the constructor's bag reads (page_size below): a standalone + # construction (benchmark/one_batch, the manual runner tests) has no + # earlier publish. + if not is_draft_worker: + set_global_server_args_for_scheduler(server_args) self.draft_attention_backend = resolve_draft_attention_backend( draft_attention_backend=draft_attention_backend, server_args=server_args, @@ -332,7 +339,7 @@ class ModelRunner: server_args.speculative_algorithm ) self.capture_tail_hooks = [] - self.page_size = server_args.page_size + self.page_size = get_schedule().page_size self.req_to_token_pool = req_to_token_pool self.token_to_kv_pool_allocator = token_to_kv_pool_allocator self.mtp_draft_device_pools = () @@ -360,11 +367,6 @@ class ModelRunner: if server_args.show_time_cost: enable_show_time_cost() - # Set the global server_args in the scheduler process (target worker - # only, so a draft init cannot clobber target-derived global state). - if not self.is_draft_worker: - set_global_server_args_for_scheduler(server_args) - misc_utils.maybe_disable_chunked_prefix_cache( use_mla_backend=self.use_mla_backend, is_draft_worker=self.is_draft_worker, diff --git a/python/sglang/srt/model_executor/model_runner_components/ngram_embedding_manager.py b/python/sglang/srt/model_executor/model_runner_components/ngram_embedding_manager.py index 795fd96c6..4d034f781 100644 --- a/python/sglang/srt/model_executor/model_runner_components/ngram_embedding_manager.py +++ b/python/sglang/srt/model_executor/model_runner_components/ngram_embedding_manager.py @@ -11,6 +11,7 @@ from sglang.kernels.ops.speculative.ngram_embedding import update_token_table from sglang.srt.configs.model_config import ModelConfig from sglang.srt.managers.schedule_batch import ForwardMode from sglang.srt.mem_cache.memory_pool import ReqToTokenPool +from sglang.srt.runtime_context import get_schedule from sglang.srt.server_args import ServerArgs if TYPE_CHECKING: @@ -50,7 +51,7 @@ class NgramEmbeddingManager: dtype=torch.int32, device=device, ) - chunked_prefill_size = server_args.chunked_prefill_size + chunked_prefill_size = get_schedule().chunked_prefill_size assert ( chunked_prefill_size is not None and chunked_prefill_size > 0 ), "Ngram embedding requires chunked prefill to be enabled (chunked_prefill_size > 0)" diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index d0c3888ae..62e7340dd 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -28,7 +28,7 @@ from sglang.srt.model_executor.forward_batch_info import ( ForwardMode, compute_position, ) -from sglang.srt.runtime_context import get_exec +from sglang.srt.runtime_context import get_exec, get_schedule from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.base_spec_worker import BaseSpecWorker from sglang.srt.speculative.dflash_info import DFlashVerifyInput @@ -189,7 +189,7 @@ class DFlashWorkerV2(BaseSpecWorker): self._target_worker = target_worker self.model_runner = target_worker.model_runner self._need_mamba_verify_commit = False - self.page_size = server_args.page_size + self.page_size = get_schedule().page_size # Normalized in arg_groups.speculative_hook.handle_speculative_decoding. self.draft_window_size: Optional[int] = ( server_args.speculative_draft_window_size diff --git a/python/sglang/srt/speculative/dspark_components/dspark_planner.py b/python/sglang/srt/speculative/dspark_components/dspark_planner.py index 2d38ab0b8..a55adbc08 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_planner.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_planner.py @@ -1139,6 +1139,6 @@ def build_sps_cost_table( return load_sps_table_from_path(sps_table_path) max_batch_tokens = max( 1, - int(server_args.max_running_requests or 1) * verify_num_draft_tokens, + int(get_schedule().max_running_requests or 1) * verify_num_draft_tokens, ) return build_uninitialized_sps_table(max_batch_tokens=max_batch_tokens) diff --git a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py index 339419741..5b59f9f62 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -18,7 +18,7 @@ from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, compute_position, ) -from sglang.srt.runtime_context import get_exec, get_parallel, get_spec +from sglang.srt.runtime_context import get_exec, get_parallel, get_schedule, get_spec from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.base_spec_worker import BaseSpecWorker from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2 @@ -88,7 +88,7 @@ class DSparkWorkerV2(BaseSpecWorker): self.nccl_port = nccl_port self._target_worker = target_worker self.model_runner = target_worker.model_runner - self.page_size = server_args.page_size + self.page_size = get_schedule().page_size self.device = target_worker.device self._draft_is_moe = draft_is_deepseek_v4(server_args=server_args) @@ -229,7 +229,7 @@ class DSparkWorkerV2(BaseSpecWorker): and is_cuda() ): self._verify_epilogue = DsparkVerifyEpilogue( - max_bs=max(server_args.cuda_graph_config.decode.bs), + max_bs=max(get_exec().graph.cuda_graph_config.decode.bs), verify_num_draft_tokens=self.verify_num_draft_tokens, device=self.device, commit_ctx=CommitInjectCtx( diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 9fb14f970..fbcabdd11 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -50,6 +50,7 @@ from sglang.srt.runtime_context import ( get_exec, get_model, get_parallel, + get_schedule, get_spec, ) from sglang.srt.server_args import ServerArgs @@ -1024,7 +1025,7 @@ class EAGLEWorkerV2(BaseSpecWorker): self.gpu_id = gpu_id self.device = server_args.device self._target_worker = target_worker - self.page_size = server_args.page_size + self.page_size = get_schedule().page_size self.speculative_algorithm = SpeculativeAlgorithm.from_string( server_args.speculative_algorithm ) diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py index e31798633..427a416f2 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py @@ -44,7 +44,7 @@ from sglang.srt.model_executor.forward_batch_info import ( ) from sglang.srt.model_executor.forward_context import ForwardContext, forward_context from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig -from sglang.srt.runtime_context import attention_backends, get_spec +from sglang.srt.runtime_context import attention_backends, get_schedule, get_spec from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.base_spec_worker import BaseSpecWorker, EagleDraftWorkerBase from sglang.srt.speculative.eagle_utils import ( @@ -109,7 +109,7 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker): self.gpu_id = gpu_id self.device = server_args.device self.target_worker = target_worker - self.page_size = server_args.page_size + self.page_size = get_schedule().page_size self.speculative_algorithm = SpeculativeAlgorithm.from_string( server_args.speculative_algorithm ) @@ -697,7 +697,7 @@ class FrozenKVMTPWorkerV2(EAGLEWorkerV2): self.gpu_id = gpu_id self.device = server_args.device self._target_worker = target_worker - self.page_size = server_args.page_size + self.page_size = get_schedule().page_size self.speculative_algorithm = SpeculativeAlgorithm.from_string( server_args.speculative_algorithm ) 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 bf449b9fd..17c9b1279 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -43,6 +43,7 @@ from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, ForwardBatch, ) +from sglang.srt.runtime_context import get_schedule from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.base_spec_worker import BaseSpecWorker, EagleDraftWorkerBase from sglang.srt.speculative.draft_utils import DraftBackendFactory @@ -933,7 +934,7 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker): self.gpu_id = gpu_id self.device = server_args.device self._target_worker = target_worker - self.page_size = server_args.page_size + self.page_size = get_schedule().page_size self.speculative_algorithm = SpeculativeAlgorithm.from_string( server_args.speculative_algorithm ) diff --git a/python/sglang/srt/speculative/ngram_worker.py b/python/sglang/srt/speculative/ngram_worker.py index 84cdecdde..7901b0021 100644 --- a/python/sglang/srt/speculative/ngram_worker.py +++ b/python/sglang/srt/speculative/ngram_worker.py @@ -15,6 +15,7 @@ from sglang.srt.managers.scheduler import GenerationBatchResult from sglang.srt.managers.tp_worker import TpModelWorker from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.srt.observability.req_time_stats import set_time_batch +from sglang.srt.runtime_context import get_schedule from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.base_spec_worker import BaseSpecWorker, EagleDraftWorkerBase from sglang.srt.speculative.cpp_ngram.ngram_corpus import NgramCorpus @@ -91,7 +92,7 @@ class NGRAMWorker(BaseSpecWorker): self._target_worker = target_worker self.model_runner = target_worker.model_runner self.tp_rank = ps.tp_rank - self.page_size = server_args.page_size + self.page_size = get_schedule().page_size self.draft_token_num: int = server_args.speculative_num_draft_tokens self.max_trie_depth: int = server_args.speculative_ngram_max_trie_depth self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens diff --git a/python/sglang/srt/speculative/standalone_worker_v2.py b/python/sglang/srt/speculative/standalone_worker_v2.py index 73275bd60..f0bfd52f5 100644 --- a/python/sglang/srt/speculative/standalone_worker_v2.py +++ b/python/sglang/srt/speculative/standalone_worker_v2.py @@ -10,6 +10,7 @@ from sglang.srt.layers.moe.utils import ( speculative_moe_backend_context, ) from sglang.srt.managers.tp_worker import TpModelWorker +from sglang.srt.runtime_context import get_schedule from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.adaptive_runtime_state import ( AdaptiveController, @@ -163,7 +164,7 @@ class StandaloneWorkerV2(EAGLEWorkerV2): self.gpu_id = gpu_id self.device = server_args.device self._target_worker = target_worker - self.page_size = server_args.page_size + self.page_size = get_schedule().page_size self.speculative_algorithm = SpeculativeAlgorithm.from_string( server_args.speculative_algorithm ) diff --git a/test/registered/spec/dspark/test_dspark_sps_table.py b/test/registered/spec/dspark/test_dspark_sps_table.py index ac13c3dc2..d0b23fe27 100644 --- a/test/registered/spec/dspark/test_dspark_sps_table.py +++ b/test/registered/spec/dspark/test_dspark_sps_table.py @@ -1,7 +1,6 @@ import tempfile import unittest from pathlib import Path -from types import SimpleNamespace from sglang.srt.speculative.dspark_components.dspark_sps import ( SpsAdditiveCostTable, @@ -143,22 +142,30 @@ class TestProfileSpsTable(CustomTestCase): self.assertEqual(table.max_batch_tokens, 256) -def _build_sps_cost_table_for(*, sps_table_path): +def _build_sps_cost_table_for(testcase, *, sps_table_path): + from sglang.srt.runtime_context import get_context, get_server_args from sglang.srt.speculative.dspark_components.dspark_planner import ( build_sps_cost_table, ) - server_args = SimpleNamespace( + # The table bound reads `max_running_requests` from the published bags, so + # the case publishes it; the table path stays on the handed record, which is + # what `build_sps_cost_table` takes. + override = get_context().override_server_args( speculative_dspark_sps_table_path=sps_table_path, max_running_requests=4, ) - return build_sps_cost_table(server_args=server_args, verify_num_draft_tokens=5) + override.install() + testcase.addCleanup(override.restore) + return build_sps_cost_table( + server_args=get_server_args(), verify_num_draft_tokens=5 + ) class TestBuildSpsCostTableContract(CustomTestCase): def test_unset_table_path_returns_flat_table(self): for sps_table_path in (None, ""): - table = _build_sps_cost_table_for(sps_table_path=sps_table_path) + table = _build_sps_cost_table_for(self, sps_table_path=sps_table_path) self.assertEqual(table.sample_batch_tokens, [1]) self.assertEqual(table.sample_steps_per_sec, [1.0]) self.assertEqual(table.max_batch_tokens, 20) @@ -168,7 +175,7 @@ class TestBuildSpsCostTableContract(CustomTestCase): with tempfile.TemporaryDirectory() as tmp: path = Path(tmp) / "sps.json" path.write_text(table.to_json(), encoding="utf-8") - loaded = _build_sps_cost_table_for(sps_table_path=str(path)) + loaded = _build_sps_cost_table_for(self, sps_table_path=str(path)) self.assertEqual(loaded.sample_batch_tokens, table.sample_batch_tokens) self.assertEqual(loaded.sample_steps_per_sec, table.sample_steps_per_sec) self.assertEqual(loaded.max_batch_tokens, table.max_batch_tokens)