config: the post-publish consumers of the supplied-instance surface read the bags

config: the speculative workers take page_size from the bags

Seven worker constructors stored `self.page_size = server_args.page_size` off
the handed record. They all run after publish and all keep a copy of a
process-level value, which is the first row of the plan doc's supplied-instance
disposition table -- so they read `get_schedule().page_size`, and a post-publish
override now reaches them like it reaches every other consumer.

The supplied-instance census named the seven pairs; the exposure ratchet in the
next member pins what remains after this batch of conversions.


config: the post-publish chunked_prefill_size consumers read the bags

Four of the ten supplied-instance `chunked_prefill_size` reads are plain
post-publish consumers -- the EPLB recorder's buffer sizing, the deep-gemm
compile warmup (five reads), the KV-cache builder's effective size, and the
ngram embedding manager's assert. All are reached from runner init, so they read
`get_schedule()`.

Two are deliberately left: `create_kt_config_from_server_args` builds a config
*from a supplied record* by name and contract, and `CanaryLaunchCapacities.from_args`
is the same shape. Converting those would change what the function is, not where
it reads -- the plan doc's disposition table says so per field.


config: the remaining post-publish graph/limit consumers read the bags

Three more of the census's supplied-instance debts are plain post-publish reads: the dspark worker's
cuda-graph decode sizes, the dspark planner's SPS table bound
(`max_running_requests`), and the LoRA manager's cuda-graph moe buffers. The
dspark worker is the clearest of them -- it already read
`get_exec().graph.cuda_graph_config.decode.bs` thirty lines below the instance
read, so the file disagreed with itself about where the same value comes from.

Left where the function's contract is "build a config from the record you are
handed" rather than "read this process's config":
`create_kt_config_from_server_args`, `DllmConfig.from_server_args`,
`CanaryLaunchCapacities.from_args`, `build_compilation_config`. Changing those
would change what the function is.


config: the runner, scheduler and offload manager take page_size from the bags

The same `self.page_size = server_args.page_size` shape as the speculative
workers, in the three remaining process-owned constructors: `ModelRunner`,
`Scheduler`, and the decode-side KV offload manager. The scheduler process
publishes before any of them run. The one path that did not is `ModelRunner`
constructed standalone -- `python -m sglang.benchmark.one_batch` and the manual
runner tests build it with no prior publish, and the constructor's own publish
sat below this read -- so that publish moves above the constructor's first bag
read instead of leaving a window where the runner half-exists unpublished.

Left where the read belongs to something else: `utils/common`'s predicates are
called only from the resolution pipeline with a `resolved_view`,
`allocation_sizing` takes the config its callers supply by contract, and
`CudaVmmFeatureTransport` is tokenizer-owned -- one per tokenizer worker, which
is the per-instance boundary.

The conversion left the offload manager parking a record it no longer
reads; the parked copy goes with the read (the constructor parameter stays
-- its hicache sizing still reads it directly).
This commit is contained in:
Cheng Wan
2026-08-15 00:39:35 -07:00
committed by GitHub
parent f2ab6e306b
commit 1ab713c334
17 changed files with 57 additions and 40 deletions
@@ -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.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.mha import get_mha_host_pool_cls
from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost 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.server_args import ServerArgs
from sglang.srt.utils.common import ceil_align from sglang.srt.utils.common import ceil_align
@@ -43,8 +44,7 @@ class DecodeKVCacheOffloadManager:
) -> None: ) -> None:
self.req_to_token_pool = req_to_token_pool self.req_to_token_pool = req_to_token_pool
self.token_to_kv_pool_allocator = token_to_kv_pool_allocator self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
self.page_size = server_args.page_size self.page_size = get_schedule().page_size
self.server_args = server_args
self.request_counter = 0 self.request_counter = 0
self.tree_cache = tree_cache self.tree_cache = tree_cache
env_stride = envs.SGLANG_HICACHE_DECODE_OFFLOAD_STRIDE.get() env_stride = envs.SGLANG_HICACHE_DECODE_OFFLOAD_STRIDE.get()
@@ -35,6 +35,7 @@ from sglang.srt.observability.metrics_collector import (
ExpertDispatchCollector, ExpertDispatchCollector,
resolve_collector_class, resolve_collector_class,
) )
from sglang.srt.runtime_context import get_schedule
from sglang.srt.server_args import ServerArgs from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import Withable, get_device, get_int_env_var from sglang.srt.utils import Withable, get_device, get_int_env_var
@@ -390,7 +391,7 @@ class _DetailSinglePassGatherer(_SinglePassGatherer):
( (
expert_location_metadata.num_layers, expert_location_metadata.num_layers,
# TODO determine the max number # TODO determine the max number
server_args.chunked_prefill_size * 8, get_schedule().chunked_prefill_size * 8,
self._TOP_K_NUM, self._TOP_K_NUM,
), ),
dtype=torch.int32, dtype=torch.int32,
@@ -16,7 +16,7 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.layers.deep_gemm_wrapper.configurer import ENABLE_JIT_DEEPGEMM 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.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.server_args import ServerArgs
from sglang.srt.utils import ceil_align, ceil_div, get_available_gpu_memory, is_musa 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) # 8192, 9008, ... 16384 (step 16)
# Totally 1024 + 1024 / 2 + 2048 / 4 + 4096 / 8 + 8192 / 16 = 3072 kernels # Totally 1024 + 1024 / 2 + 2048 / 4 + 4096 / 8 + 8192 / 16 = 3072 kernels
next_m, sample_step = 1024, 2 next_m, sample_step = 1024, 2
chunked_prefill_size = get_schedule().chunked_prefill_size
max_prefill_bs = ( max_prefill_bs = (
min(server_args.chunked_prefill_size, 32 * 1024) min(chunked_prefill_size, 32 * 1024)
if server_args.chunked_prefill_size >= 1 if chunked_prefill_size >= 1
else 16 * 1024 else 16 * 1024
) )
while next_m < max_prefill_bs: while next_m < max_prefill_bs:
@@ -81,10 +82,11 @@ def update_deep_gemm_config(gpu_id: int, server_args: ServerArgs):
else: else:
# When fast warmup isn't enabled, generate m_max and compile all the covered Ms. # When fast warmup isn't enabled, generate m_max and compile all the covered Ms.
m_max = 1024 * 16 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 m_max = 1024 * 64
elif server_args.chunked_prefill_size > 8192: elif chunked_prefill_size > 8192:
m_max = server_args.chunked_prefill_size * 2 m_max = chunked_prefill_size * 2
m_max = min(1024 * 128, m_max) m_max = min(1024 * 128, m_max)
_BUILTIN_M_LIST += list(range(1, m_max + 1)) _BUILTIN_M_LIST += list(range(1, m_max + 1))
+2 -2
View File
@@ -46,7 +46,7 @@ from sglang.srt.lora.utils import (
) )
from sglang.srt.managers.io_struct import LoRAUpdateOutput from sglang.srt.managers.io_struct import LoRAUpdateOutput
from sglang.srt.model_executor.forward_batch_info import ForwardBatch 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.server_args import ServerArgs
from sglang.srt.utils import get_available_gpu_memory, replace_submodule from sglang.srt.utils import get_available_gpu_memory, replace_submodule
from sglang.srt.utils.hf_transformers_utils import AutoConfig 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 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 max_loras = server_args.max_loras_per_batch
for module in model.modules(): for module in model.modules():
if isinstance(module, FusedMoEWithLoRA): if isinstance(module, FusedMoEWithLoRA):
+1 -1
View File
@@ -436,7 +436,7 @@ class Scheduler(
self.spec_algorithm = SpeculativeAlgorithm.from_string( self.spec_algorithm = SpeculativeAlgorithm.from_string(
server_args.speculative_algorithm 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_hierarchical_cache = server_args.enable_hierarchical_cache
self.enable_session_radix_cache = server_args.enable_session_radix_cache self.enable_session_radix_cache = server_args.enable_session_radix_cache
self.enable_hicache_storage = server_args.hicache_storage_backend is not None self.enable_hicache_storage = server_args.hicache_storage_backend is not None
@@ -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.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.registry import TreeCacheBuildContext, create_tree_cache 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.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: if TYPE_CHECKING:
@@ -205,7 +205,7 @@ def build_kv_cache(
"with Mamba/SSM models" "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: if model_config.is_multimodal and uses_transformers_backend:
effective_chunked_prefill_size = None effective_chunked_prefill_size = None
@@ -313,6 +313,13 @@ class ModelRunner:
self.dist_port = nccl_port self.dist_port = nccl_port
self.server_args = server_args self.server_args = server_args
self.is_draft_worker = is_draft_worker 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( self.draft_attention_backend = resolve_draft_attention_backend(
draft_attention_backend=draft_attention_backend, draft_attention_backend=draft_attention_backend,
server_args=server_args, server_args=server_args,
@@ -332,7 +339,7 @@ class ModelRunner:
server_args.speculative_algorithm server_args.speculative_algorithm
) )
self.capture_tail_hooks = [] 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.req_to_token_pool = req_to_token_pool
self.token_to_kv_pool_allocator = token_to_kv_pool_allocator self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
self.mtp_draft_device_pools = () self.mtp_draft_device_pools = ()
@@ -360,11 +367,6 @@ class ModelRunner:
if server_args.show_time_cost: if server_args.show_time_cost:
enable_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( misc_utils.maybe_disable_chunked_prefix_cache(
use_mla_backend=self.use_mla_backend, use_mla_backend=self.use_mla_backend,
is_draft_worker=self.is_draft_worker, is_draft_worker=self.is_draft_worker,
@@ -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.configs.model_config import ModelConfig
from sglang.srt.managers.schedule_batch import ForwardMode from sglang.srt.managers.schedule_batch import ForwardMode
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.runtime_context import get_schedule
from sglang.srt.server_args import ServerArgs from sglang.srt.server_args import ServerArgs
if TYPE_CHECKING: if TYPE_CHECKING:
@@ -50,7 +51,7 @@ class NgramEmbeddingManager:
dtype=torch.int32, dtype=torch.int32,
device=device, device=device,
) )
chunked_prefill_size = server_args.chunked_prefill_size chunked_prefill_size = get_schedule().chunked_prefill_size
assert ( assert (
chunked_prefill_size is not None and chunked_prefill_size > 0 chunked_prefill_size is not None and chunked_prefill_size > 0
), "Ngram embedding requires chunked prefill to be enabled (chunked_prefill_size > 0)" ), "Ngram embedding requires chunked prefill to be enabled (chunked_prefill_size > 0)"
@@ -28,7 +28,7 @@ from sglang.srt.model_executor.forward_batch_info import (
ForwardMode, ForwardMode,
compute_position, 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.server_args import ServerArgs
from sglang.srt.speculative.base_spec_worker import BaseSpecWorker from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
from sglang.srt.speculative.dflash_info import DFlashVerifyInput from sglang.srt.speculative.dflash_info import DFlashVerifyInput
@@ -189,7 +189,7 @@ class DFlashWorkerV2(BaseSpecWorker):
self._target_worker = target_worker self._target_worker = target_worker
self.model_runner = target_worker.model_runner self.model_runner = target_worker.model_runner
self._need_mamba_verify_commit = False 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. # Normalized in arg_groups.speculative_hook.handle_speculative_decoding.
self.draft_window_size: Optional[int] = ( self.draft_window_size: Optional[int] = (
server_args.speculative_draft_window_size server_args.speculative_draft_window_size
@@ -1139,6 +1139,6 @@ def build_sps_cost_table(
return load_sps_table_from_path(sps_table_path) return load_sps_table_from_path(sps_table_path)
max_batch_tokens = max( max_batch_tokens = max(
1, 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) return build_uninitialized_sps_table(max_batch_tokens=max_batch_tokens)
@@ -18,7 +18,7 @@ from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode, CaptureHiddenMode,
compute_position, 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.server_args import ServerArgs
from sglang.srt.speculative.base_spec_worker import BaseSpecWorker from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2 from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2
@@ -88,7 +88,7 @@ class DSparkWorkerV2(BaseSpecWorker):
self.nccl_port = nccl_port self.nccl_port = nccl_port
self._target_worker = target_worker self._target_worker = target_worker
self.model_runner = target_worker.model_runner 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.device = target_worker.device
self._draft_is_moe = draft_is_deepseek_v4(server_args=server_args) self._draft_is_moe = draft_is_deepseek_v4(server_args=server_args)
@@ -229,7 +229,7 @@ class DSparkWorkerV2(BaseSpecWorker):
and is_cuda() and is_cuda()
): ):
self._verify_epilogue = DsparkVerifyEpilogue( 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, verify_num_draft_tokens=self.verify_num_draft_tokens,
device=self.device, device=self.device,
commit_ctx=CommitInjectCtx( commit_ctx=CommitInjectCtx(
@@ -50,6 +50,7 @@ from sglang.srt.runtime_context import (
get_exec, get_exec,
get_model, get_model,
get_parallel, get_parallel,
get_schedule,
get_spec, get_spec,
) )
from sglang.srt.server_args import ServerArgs from sglang.srt.server_args import ServerArgs
@@ -1024,7 +1025,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
self.gpu_id = gpu_id self.gpu_id = gpu_id
self.device = server_args.device self.device = server_args.device
self._target_worker = target_worker 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( self.speculative_algorithm = SpeculativeAlgorithm.from_string(
server_args.speculative_algorithm server_args.speculative_algorithm
) )
@@ -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.forward_context import ForwardContext, forward_context
from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig 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.server_args import ServerArgs
from sglang.srt.speculative.base_spec_worker import BaseSpecWorker, EagleDraftWorkerBase from sglang.srt.speculative.base_spec_worker import BaseSpecWorker, EagleDraftWorkerBase
from sglang.srt.speculative.eagle_utils import ( from sglang.srt.speculative.eagle_utils import (
@@ -109,7 +109,7 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker):
self.gpu_id = gpu_id self.gpu_id = gpu_id
self.device = server_args.device self.device = server_args.device
self.target_worker = target_worker 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( self.speculative_algorithm = SpeculativeAlgorithm.from_string(
server_args.speculative_algorithm server_args.speculative_algorithm
) )
@@ -697,7 +697,7 @@ class FrozenKVMTPWorkerV2(EAGLEWorkerV2):
self.gpu_id = gpu_id self.gpu_id = gpu_id
self.device = server_args.device self.device = server_args.device
self._target_worker = target_worker 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( self.speculative_algorithm = SpeculativeAlgorithm.from_string(
server_args.speculative_algorithm server_args.speculative_algorithm
) )
@@ -43,6 +43,7 @@ from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode, CaptureHiddenMode,
ForwardBatch, ForwardBatch,
) )
from sglang.srt.runtime_context import get_schedule
from sglang.srt.server_args import ServerArgs from sglang.srt.server_args import ServerArgs
from sglang.srt.speculative.base_spec_worker import BaseSpecWorker, EagleDraftWorkerBase from sglang.srt.speculative.base_spec_worker import BaseSpecWorker, EagleDraftWorkerBase
from sglang.srt.speculative.draft_utils import DraftBackendFactory from sglang.srt.speculative.draft_utils import DraftBackendFactory
@@ -933,7 +934,7 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
self.gpu_id = gpu_id self.gpu_id = gpu_id
self.device = server_args.device self.device = server_args.device
self._target_worker = target_worker 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( self.speculative_algorithm = SpeculativeAlgorithm.from_string(
server_args.speculative_algorithm server_args.speculative_algorithm
) )
@@ -15,6 +15,7 @@ from sglang.srt.managers.scheduler import GenerationBatchResult
from sglang.srt.managers.tp_worker import TpModelWorker from sglang.srt.managers.tp_worker import TpModelWorker
from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.observability.req_time_stats import set_time_batch 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.server_args import ServerArgs
from sglang.srt.speculative.base_spec_worker import BaseSpecWorker, EagleDraftWorkerBase from sglang.srt.speculative.base_spec_worker import BaseSpecWorker, EagleDraftWorkerBase
from sglang.srt.speculative.cpp_ngram.ngram_corpus import NgramCorpus from sglang.srt.speculative.cpp_ngram.ngram_corpus import NgramCorpus
@@ -91,7 +92,7 @@ class NGRAMWorker(BaseSpecWorker):
self._target_worker = target_worker self._target_worker = target_worker
self.model_runner = target_worker.model_runner self.model_runner = target_worker.model_runner
self.tp_rank = ps.tp_rank 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.draft_token_num: int = server_args.speculative_num_draft_tokens
self.max_trie_depth: int = server_args.speculative_ngram_max_trie_depth self.max_trie_depth: int = server_args.speculative_ngram_max_trie_depth
self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens
@@ -10,6 +10,7 @@ from sglang.srt.layers.moe.utils import (
speculative_moe_backend_context, speculative_moe_backend_context,
) )
from sglang.srt.managers.tp_worker import TpModelWorker 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.server_args import ServerArgs
from sglang.srt.speculative.adaptive_runtime_state import ( from sglang.srt.speculative.adaptive_runtime_state import (
AdaptiveController, AdaptiveController,
@@ -163,7 +164,7 @@ class StandaloneWorkerV2(EAGLEWorkerV2):
self.gpu_id = gpu_id self.gpu_id = gpu_id
self.device = server_args.device self.device = server_args.device
self._target_worker = target_worker 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( self.speculative_algorithm = SpeculativeAlgorithm.from_string(
server_args.speculative_algorithm server_args.speculative_algorithm
) )
@@ -1,7 +1,6 @@
import tempfile import tempfile
import unittest import unittest
from pathlib import Path from pathlib import Path
from types import SimpleNamespace
from sglang.srt.speculative.dspark_components.dspark_sps import ( from sglang.srt.speculative.dspark_components.dspark_sps import (
SpsAdditiveCostTable, SpsAdditiveCostTable,
@@ -143,22 +142,30 @@ class TestProfileSpsTable(CustomTestCase):
self.assertEqual(table.max_batch_tokens, 256) 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 ( from sglang.srt.speculative.dspark_components.dspark_planner import (
build_sps_cost_table, 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, speculative_dspark_sps_table_path=sps_table_path,
max_running_requests=4, 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): class TestBuildSpsCostTableContract(CustomTestCase):
def test_unset_table_path_returns_flat_table(self): def test_unset_table_path_returns_flat_table(self):
for sps_table_path in (None, ""): 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_batch_tokens, [1])
self.assertEqual(table.sample_steps_per_sec, [1.0]) self.assertEqual(table.sample_steps_per_sec, [1.0])
self.assertEqual(table.max_batch_tokens, 20) self.assertEqual(table.max_batch_tokens, 20)
@@ -168,7 +175,7 @@ class TestBuildSpsCostTableContract(CustomTestCase):
with tempfile.TemporaryDirectory() as tmp: with tempfile.TemporaryDirectory() as tmp:
path = Path(tmp) / "sps.json" path = Path(tmp) / "sps.json"
path.write_text(table.to_json(), encoding="utf-8") 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_batch_tokens, table.sample_batch_tokens)
self.assertEqual(loaded.sample_steps_per_sec, table.sample_steps_per_sec) self.assertEqual(loaded.sample_steps_per_sec, table.sample_steps_per_sec)
self.assertEqual(loaded.max_batch_tokens, table.max_batch_tokens) self.assertEqual(loaded.max_batch_tokens, table.max_batch_tokens)