chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Clean up __init__ function of the scheduler and event loop for PD (#15298)

Browse Source
This commit is contained in:
Lianmin Zheng
2025-12-18 01:35:14 -08:00
committed by GitHub
parent 4792d1f452
commit d1f0063262
9 changed files with 662 additions and 641 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler.py
+371 -343
View File
@@ -246,6 +246,7 @@ class Scheduler(
self.moe_ep_size = server_args.ep_size
self.pp_size = server_args.pp_size
self.dp_size = server_args.dp_size
self.nccl_port = port_args.nccl_port
self.schedule_policy = server_args.schedule_policy
self.enable_priority_scheduling = server_args.enable_priority_scheduling
self.abort_on_priority_when_disabled = (
@@ -292,9 +293,11 @@ class Scheduler(
# Init model config
self.model_config = ModelConfig.from_server_args(server_args)
# Init diffusion LLM config
self.dllm_config = DllmConfig.from_server_args(server_args)
self.dllm_config = ( # For diffusion LLM
DllmConfig.from_server_args(server_args)
if server_args.dllm_algorithm is not None
else None
)
# Init metrics stats
self.init_metrics(tp_rank, pp_rank, dp_rank)
@@ -309,348 +312,49 @@ class Scheduler(
# Init tokenizer
self.init_tokenizer()
# Init moe config
self.init_moe_config()
# Init GEMM config (FP8 GEMM, etc.)
self.init_gemm_config()
# Check whether overlap can be enabled
if not self.is_generation:
self.enable_overlap = False
logger.info("Overlap scheduler is disabled for embedding models.")
# Init moe config and GEMM config (FP8 GEMM, etc.)
self.init_moe_gemm_config()
# Launch a tensor parallel worker
from sglang.srt.managers.tp_worker import TpModelWorker
self.tp_worker = TpModelWorker(
server_args=server_args,
gpu_id=gpu_id,
tp_rank=tp_rank,
moe_ep_rank=moe_ep_rank,
pp_rank=pp_rank,
dp_rank=dp_rank,
nccl_port=port_args.nccl_port,
)
# Launch a draft worker for speculative decoding
draft_worker_kwargs = dict(
gpu_id=gpu_id,
tp_rank=tp_rank,
moe_ep_rank=moe_ep_rank,
server_args=server_args,
nccl_port=port_args.nccl_port,
target_worker=self.tp_worker,
dp_rank=dp_rank,
)
if server_args.speculative_draft_load_format is not None:
server_args.load_format = server_args.speculative_draft_load_format
logger.info(
f"Using draft model load_format: '{server_args.speculative_draft_load_format}'"
)
# Draft workers are looked up via `SpeculativeAlgorithm` registry; new
# algorithms should register their factory instead of patching this code.
if self.spec_algorithm.is_eagle():
draft_worker_kwargs["enable_overlap"] = self.enable_overlap
self.draft_worker = self.spec_algorithm.create_draft_worker(
**draft_worker_kwargs
)
# Dispatch the model worker
if self.spec_algorithm.is_none():
self.model_worker = self.tp_worker
else:
self.model_worker = self.draft_worker
# Get token and memory info from the model worker
(
self.max_total_num_tokens,
self.max_prefill_tokens,
self.max_running_requests,
self.max_queued_requests,
self.max_req_len,
self.max_req_input_len,
self.random_seed,
self.device,
_,
_,
_,
) = self.tp_worker.get_worker_info()
if get_global_server_args().pp_max_micro_batch_size is None:
get_global_server_args().pp_max_micro_batch_size = max(
self.max_running_requests // server_args.pp_size, 1
)
self.tp_group = self.tp_worker.get_tp_group()
self.tp_cpu_group = self.tp_group.cpu_group
self.attn_tp_group = self.tp_worker.get_attention_tp_group()
self.attn_tp_cpu_group = self.tp_worker.get_attention_tp_cpu_group()
self.pp_group = get_pp_group()
self.world_group = get_world_group()
# With DP attention enabled, the entry rank is attn_tp_rank==0;
# otherwise the entry rank is TP group local rank 0.
# For #11910, use the CPU communication group to broadcast VLM Python objects,
# avoiding any coupling with CUDA streams/devices.
if self.server_args.enable_dp_attention:
self.cpu_group = self.attn_tp_cpu_group
self.entry_rank = self.attn_tp_group.first_rank
self.is_entry_rank = self.attn_tp_rank == 0
else:
self.cpu_group = self.tp_cpu_group
self.entry_rank = self.tp_group.first_rank
self.is_entry_rank = self.tp_group.rank_in_group == 0
self.pad_input_ids_func = self.tp_worker.get_pad_input_ids_func()
set_random_seed(self.random_seed)
# Hybrid memory pool
self.is_hybrid_swa = self.tp_worker.is_hybrid_swa
self.is_hybrid_ssm = (
self.tp_worker.model_runner.hybrid_gdn_config is not None
or self.tp_worker.model_runner.mamba2_config is not None
)
if self.is_hybrid_swa:
self.sliding_window_size = self.tp_worker.sliding_window_size
self.full_tokens_per_layer, self.swa_tokens_per_layer = (
self.tp_worker.get_tokens_per_layer_info()
)
# Print debug info
if tp_rank == 0:
avail_mem = get_available_gpu_memory(
self.device, self.gpu_id, empty_cache=False
)
logger.info(
f"max_total_num_tokens={self.max_total_num_tokens}, "
f"chunked_prefill_size={server_args.chunked_prefill_size}, "
f"max_prefill_tokens={self.max_prefill_tokens}, "
f"max_running_requests={self.max_running_requests}, "
f"context_len={self.model_config.context_len}, "
f"{'available_cpu_mem' if self.device == 'cpu' else 'available_gpu_mem'}={avail_mem:.2f} GB"
)
self.init_model_worker()
# Init cache using the existing memory pool
self.init_cache_with_memory_pool()
# Init running status
self.waiting_queue: List[Req] = []
# The running decoding batch for continuous batching
self.running_batch: ScheduleBatch = ScheduleBatch(reqs=[], batch_is_full=False)
# The current forward batch
self.cur_batch: Optional[ScheduleBatch] = None
# The current split prefill batch
self.split_prefill_batch: Optional[ScheduleBatch] = None
# The last forward batch
self.last_batch: Optional[ScheduleBatch] = None
self.forward_ct = 0
self.last_prefill_tokens = 0
self.return_health_check_ct = 0
self.num_retracted_reqs: int = 0
self.num_paused_reqs: int = 0
self.sessions: Dict[str, Session] = {}
self.forward_sleep_time = None
self._engine_paused = False
self.init_running_status()
# Init chunked prefill
self.chunked_prefill_size = server_args.chunked_prefill_size
if self.dllm_config is not None:
# We currently leverage chunked prefill to implement block diffusion
# for diffusion LLM.
self.chunked_prefill_size = self.dllm_config.block_size
if self.chunked_prefill_size <= 0: # -1 means disable
self.chunked_prefill_size = None
self.chunked_req = None
self.is_mixed_chunk = (
self.chunked_prefill_size is not None and server_args.enable_mixed_chunk
)
self.enable_dynamic_chunking = (
server_args.enable_dynamic_chunking and self.pp_size > 1
)
# Init the dynamic chunking predictor for PP
if self.enable_dynamic_chunking:
try:
self.profile_and_init_predictor()
except Exception as e:
logger.warning(
f"[PP Dynamic Chunk] Failed to profile prefill latency: {e}. "
"Dynamic chunking will be disabled."
)
self.enable_dynamic_chunking = False
self.init_chunked_prefill()
# Init the grammar backend for constrained generation
self.grammar_queue: List[Req] = []
if not server_args.skip_tokenizer_init:
self.grammar_backend = create_grammar_backend(
server_args,
self.tokenizer,
self.model_config.vocab_size,
self.model_config.hf_eos_token_id,
)
else:
self.grammar_backend = None
self.init_grammar_backend()
# Init schedule policy and new token estimation
self.policy = SchedulePolicy(
self.schedule_policy,
self.tree_cache,
self.enable_hierarchical_cache,
self.enable_priority_scheduling,
self.schedule_low_priority_values_first,
)
self.schedule_enhancer = None
if SCHEDULER_DECREASE_PREFILL_IDLE:
self.schedule_enhancer = SchedulerEnhancer(
self.dp_size,
self.attn_tp_size,
self.tp_worker,
self.max_running_requests,
server_args,
)
# Enable preemption for priority scheduling.
self.try_preemption = self.enable_priority_scheduling
self.init_new_token_ratio = min(
envs.SGLANG_INIT_NEW_TOKEN_RATIO.get()
* server_args.schedule_conservativeness,
1.0,
)
self.min_new_token_ratio = min(
self.init_new_token_ratio * envs.SGLANG_MIN_NEW_TOKEN_RATIO_FACTOR.get(),
1.0,
)
self.new_token_ratio_decay = (
self.init_new_token_ratio - self.min_new_token_ratio
) / envs.SGLANG_NEW_TOKEN_RATIO_DECAY_STEPS.get()
self.new_token_ratio = self.init_new_token_ratio
self.init_schedule_policy()
# Init watchdog thread
self.watchdog = SchedulerWatchdog(
self, watchdog_timeout=server_args.watchdog_timeout
)
if (x := server_args.soft_watchdog_timeout) is not None:
self.soft_watchdog = SchedulerWatchdog(self, watchdog_timeout=x, soft=True)
# Init watchdog, memory saver, input blocker and recv skipper
self.init_watch_dog_memory_saver_input_blocker()
# Init memory saver, profiler and metric stats
self.memory_saver_adapter = TorchMemorySaverAdapter.create(
enable=server_args.enable_memory_saver
)
self.offload_tags = set()
# Init profiler
self.init_profiler()
self.recv_skipper = SchedulerRecvSkipper.maybe_create(server_args)
self.input_blocker = (
SchedulerInputBlocker(noop=self.attn_tp_rank != 0)
if get_bool_env_var("SGLANG_ENABLE_COLOCATED_BATCH_GEN")
else None
)
# Init disaggregation
self.init_disaggregation()
if self.enable_kv_cache_events:
self.init_kv_events(server_args.kv_events_config)
if envs.SGLANG_LOG_GC.get():
configure_gc_logger()
# Init overlap
self.init_overlap()
# Init prefill kv split size when deterministic inference is enabled with various attention backends
self.init_deterministic_inference_config()
# Init overlap
self.init_overlap()
# Init mlp sync flag
self.require_mlp_sync = require_mlp_sync(server_args)
if (
self.server_args.language_only
and self.server_args.encoder_transfer_backend == "zmq_to_scheduler"
):
self.mm_receiver = MMReceiver(
server_args,
hf_config=self.model_config.hf_config,
tp_rank=self.tp_rank,
pp_rank=self.pp_rank,
)
# Init request dispatcher
self._request_dispatcher = TypeBasedDispatcher(
[
(TokenizedGenerateReqInput, self.handle_generate_request),
(TokenizedEmbeddingReqInput, self.handle_embedding_request),
(BatchTokenizedGenerateReqInput, self.handle_batch_generate_request),
(BatchTokenizedEmbeddingReqInput, self.handle_batch_embedding_request),
(FlushCacheReqInput, self.flush_cache_wrapped),
(ClearHiCacheReqInput, self.clear_hicache_storage_wrapped),
(AbortReq, self.abort_request),
(OpenSessionReqInput, self.open_session),
(CloseSessionReqInput, self.close_session),
(UpdateWeightFromDiskReqInput, self.update_weights_from_disk),
(InitWeightsUpdateGroupReqInput, self.init_weights_update_group),
(DestroyWeightsUpdateGroupReqInput, self.destroy_weights_update_group),
(
InitWeightsSendGroupForRemoteInstanceReqInput,
self.init_weights_send_group_for_remote_instance,
),
(
SendWeightsToRemoteInstanceReqInput,
self.send_weights_to_remote_instance,
),
(
UpdateWeightsFromDistributedReqInput,
self.update_weights_from_distributed,
),
(UpdateWeightsFromTensorReqInput, self.update_weights_from_tensor),
(UpdateWeightsFromIPCReqInput, self.update_weights_from_ipc),
(GetWeightsByNameReqInput, self.get_weights_by_name),
(ReleaseMemoryOccupationReqInput, self.release_memory_occupation),
(ResumeMemoryOccupationReqInput, self.resume_memory_occupation),
(CheckWeightsReqInput, self.check_weights),
(SlowDownReqInput, self.slow_down),
(ProfileReq, self.profile),
(FreezeGCReq, self.handle_freeze_gc),
(GetInternalStateReq, self.get_internal_state),
(SetInternalStateReq, self.set_internal_state),
(RpcReqInput, self.handle_rpc_request),
(ExpertDistributionReq, self.expert_distribution_handle),
(LoadLoRAAdapterReqInput, self.load_lora_adapter),
(UnloadLoRAAdapterReqInput, self.unload_lora_adapter),
(GetLoadReqInput, self.get_load),
(PauseGenerationReqInput, self.pause_generation),
(ContinueGenerationReqInput, self.continue_generation),
]
)
self.init_request_dispatcher()
def init_sockets(self, server_args: ServerArgs, port_args: PortArgs):
context = zmq.Context(2)
self.idle_sleeper = None
class SenderWrapper:
def __init__(self, socket: zmq.Socket):
self.socket = socket
def send_output(
self,
output: Union[BaseReq, BaseBatchReq],
recv_obj: Optional[Union[BaseReq, BaseBatchReq]] = None,
):
if self.socket is None:
return
if (
isinstance(recv_obj, BaseReq)
and recv_obj.http_worker_ipc is not None
and output.http_worker_ipc is None
):
# handle communicator reqs for multi-http worker case
output.http_worker_ipc = recv_obj.http_worker_ipc
self.socket.send_pyobj(output)
if self.pp_rank == 0 and self.attn_tp_rank == 0:
self.recv_from_tokenizer = get_zmq_socket(
context, zmq.PULL, port_args.scheduler_input_ipc_name, False
@@ -694,23 +398,6 @@ class Scheduler(
context, zmq.PUSH, port_args.metrics_ipc_name, False
)
def init_deterministic_inference_config(self):
"""Initialize deterministic inference configuration for different attention backends."""
if not self.server_args.enable_deterministic_inference:
self.truncation_align_size = None
return
backend_sizes = {
"flashinfer": ("SGLANG_FLASHINFER_PREFILL_SPLIT_TILE_SIZE", 4096),
"triton": ("SGLANG_TRITON_PREFILL_TRUNCATION_ALIGN_SIZE", 4096),
}
env_var, default_size = backend_sizes.get(
self.server_args.attention_backend, (None, None)
)
self.truncation_align_size = (
get_int_env_var(env_var, default_size) if env_var else None
)
def init_tokenizer(self):
server_args = self.server_args
self.is_generation = self.model_config.is_generation
@@ -744,13 +431,141 @@ class Scheduler(
reasoning_parser.detector.think_end_token, add_special_tokens=False
)[0]
def init_moe_gemm_config(self):
if hasattr(self.model_config.hf_config, "num_experts_per_tok"):
initialize_moe_config(self.server_args)
# Initialize GEMM-related configuration (currently FP8 Blockwise GEMM backend).
# Other GEMM backends (e.g. FP4, BF16, etc.) can be added here in the future.
# This is needed for FP8 quantization.
initialize_fp8_gemm_config(self.server_args)
# This must be called after initialize_moe_config
self.require_mlp_sync = require_mlp_sync(self.server_args)
def init_model_worker(self):
from sglang.srt.managers.tp_worker import TpModelWorker
self.tp_worker = TpModelWorker(
server_args=self.server_args,
gpu_id=self.gpu_id,
tp_rank=self.tp_rank,
moe_ep_rank=self.moe_ep_rank,
pp_rank=self.pp_rank,
dp_rank=self.dp_rank,
nccl_port=self.nccl_port,
)
# Launch a draft worker for speculative decoding
draft_worker_kwargs = dict(
server_args=self.server_args,
gpu_id=self.gpu_id,
tp_rank=self.tp_rank,
moe_ep_rank=self.moe_ep_rank,
nccl_port=self.nccl_port,
target_worker=self.tp_worker,
dp_rank=self.dp_rank,
)
if self.server_args.speculative_draft_load_format is not None:
self.server_args.load_format = (
self.server_args.speculative_draft_load_format
)
logger.info(
f"Using draft model load_format: '{self.server_args.speculative_draft_load_format}'"
)
# Draft workers are looked up via `SpeculativeAlgorithm` registry; new
# algorithms should register their factory instead of patching this code.
if self.spec_algorithm.is_eagle():
draft_worker_kwargs["enable_overlap"] = self.enable_overlap
self.draft_worker = self.spec_algorithm.create_draft_worker(
**draft_worker_kwargs
)
# Dispatch the model worker
if self.spec_algorithm.is_none():
self.model_worker = self.tp_worker
else:
self.model_worker = self.draft_worker
# Get token and memory info from the model worker
(
self.max_total_num_tokens,
self.max_prefill_tokens,
self.max_running_requests,
self.max_queued_requests,
self.max_req_len,
self.max_req_input_len,
self.random_seed,
self.device,
_,
_,
_,
) = self.tp_worker.get_worker_info()
if get_global_server_args().pp_max_micro_batch_size is None:
get_global_server_args().pp_max_micro_batch_size = max(
self.max_running_requests // self.pp_size, 1
)
self.tp_group = self.tp_worker.get_tp_group()
self.tp_cpu_group = self.tp_group.cpu_group
self.attn_tp_group = self.tp_worker.get_attention_tp_group()
self.attn_tp_cpu_group = self.tp_worker.get_attention_tp_cpu_group()
self.pp_group = get_pp_group()
self.world_group = get_world_group()
# With DP attention enabled, the entry rank is attn_tp_rank==0;
# otherwise the entry rank is TP group local rank 0.
# For #11910, use the CPU communication group to broadcast VLM Python objects,
# avoiding any coupling with CUDA streams/devices.
if self.server_args.enable_dp_attention:
self.cpu_group = self.attn_tp_cpu_group
self.entry_rank = self.attn_tp_group.first_rank
self.is_entry_rank = self.attn_tp_rank == 0
else:
self.cpu_group = self.tp_cpu_group
self.entry_rank = self.tp_group.first_rank
self.is_entry_rank = self.tp_group.rank_in_group == 0
self.pad_input_ids_func = self.tp_worker.get_pad_input_ids_func()
set_random_seed(self.random_seed)
# Print debug info
if self.tp_rank == 0:
avail_mem = get_available_gpu_memory(
self.device, self.gpu_id, empty_cache=False
)
logger.info(
f"max_total_num_tokens={self.max_total_num_tokens}, "
f"chunked_prefill_size={self.server_args.chunked_prefill_size}, "
f"max_prefill_tokens={self.max_prefill_tokens}, "
f"max_running_requests={self.max_running_requests}, "
f"context_len={self.model_config.context_len}, "
f"{'available_cpu_mem' if self.device == 'cpu' else 'available_gpu_mem'}={avail_mem:.2f} GB"
)
def init_cache_with_memory_pool(self):
server_args = self.server_args
# Hybrid memory pool configs
self.is_hybrid_swa = self.tp_worker.is_hybrid_swa
self.is_hybrid_ssm = (
self.tp_worker.model_runner.hybrid_gdn_config is not None
or self.tp_worker.model_runner.mamba2_config is not None
)
if self.is_hybrid_swa:
self.sliding_window_size = self.tp_worker.sliding_window_size
self.full_tokens_per_layer, self.swa_tokens_per_layer = (
self.tp_worker.get_tokens_per_layer_info()
)
self.req_to_token_pool, self.token_to_kv_pool_allocator = (
self.tp_worker.get_memory_pool()
)
# Create cache
params = CacheInitParams(
disable=server_args.disable_radix_cache,
req_to_token_pool=self.req_to_token_pool,
@@ -850,6 +665,126 @@ class Scheduler(
embedding_cache_size = envs.SGLANG_VLM_CACHE_SIZE_MB.get()
init_mm_embedding_cache(embedding_cache_size * 1024 * 1024)
def init_running_status(self):
self.waiting_queue: List[Req] = []
# The running decoding batch for continuous batching
self.running_batch: ScheduleBatch = ScheduleBatch(reqs=[], batch_is_full=False)
# The current forward batch
self.cur_batch: Optional[ScheduleBatch] = None
# The current split prefill batch
self.split_prefill_batch: Optional[ScheduleBatch] = None
# The last forward batch
self.last_batch: Optional[ScheduleBatch] = None
self.forward_ct = 0
self.last_prefill_tokens = 0
self.return_health_check_ct = 0
self.num_retracted_reqs: int = 0
self.num_paused_reqs: int = 0
self.sessions: Dict[str, Session] = {}
self.forward_sleep_time = None
self._engine_paused = False
def init_chunked_prefill(self):
# Init chunked prefill
self.chunked_prefill_size = self.server_args.chunked_prefill_size
if self.dllm_config is not None:
# We currently leverage chunked prefill to implement block diffusion
# for diffusion LLM.
self.chunked_prefill_size = self.dllm_config.block_size
if self.chunked_prefill_size <= 0: # -1 means disable
self.chunked_prefill_size = None
self.chunked_req = None
self.is_mixed_chunk = (
self.chunked_prefill_size is not None
and self.server_args.enable_mixed_chunk
)
# Init the dynamic chunking predictor for PP
self.enable_dynamic_chunking = (
self.server_args.enable_dynamic_chunking and self.pp_size > 1
)
if self.enable_dynamic_chunking:
try:
self.profile_and_init_predictor()
except Exception as e:
logger.warning(
f"[PP Dynamic Chunk] Failed to profile prefill latency: {e}. "
"Dynamic chunking will be disabled."
)
self.enable_dynamic_chunking = False
def init_grammar_backend(self):
self.grammar_queue: List[Req] = []
if not self.server_args.skip_tokenizer_init:
self.grammar_backend = create_grammar_backend(
self.server_args,
self.tokenizer,
self.model_config.vocab_size,
self.model_config.hf_eos_token_id,
)
else:
self.grammar_backend = None
def init_schedule_policy(self):
# Init schedule policy and new token estimation
self.policy = SchedulePolicy(
self.schedule_policy,
self.tree_cache,
self.enable_hierarchical_cache,
self.enable_priority_scheduling,
self.schedule_low_priority_values_first,
)
self.schedule_enhancer = None
if SCHEDULER_DECREASE_PREFILL_IDLE:
self.schedule_enhancer = SchedulerEnhancer(
self.dp_size,
self.attn_tp_size,
self.tp_worker,
self.max_running_requests,
self.server_args,
)
# Enable preemption for priority scheduling.
self.try_preemption = self.enable_priority_scheduling
self.init_new_token_ratio = min(
envs.SGLANG_INIT_NEW_TOKEN_RATIO.get()
* self.server_args.schedule_conservativeness,
1.0,
)
self.min_new_token_ratio = min(
self.init_new_token_ratio * envs.SGLANG_MIN_NEW_TOKEN_RATIO_FACTOR.get(),
1.0,
)
self.new_token_ratio_decay = (
self.init_new_token_ratio - self.min_new_token_ratio
) / envs.SGLANG_NEW_TOKEN_RATIO_DECAY_STEPS.get()
self.new_token_ratio = self.init_new_token_ratio
def init_watch_dog_memory_saver_input_blocker(self):
# Start watchdog thread
self.watchdog = SchedulerWatchdog(
self, watchdog_timeout=self.server_args.watchdog_timeout
)
if (x := self.server_args.soft_watchdog_timeout) is not None:
self.soft_watchdog = SchedulerWatchdog(self, watchdog_timeout=x, soft=True)
# Init memory saver, profiler and metric stats
self.memory_saver_adapter = TorchMemorySaverAdapter.create(
enable=self.server_args.enable_memory_saver
)
self.offload_tags = set()
# Init recv skipper and input blocker
self.recv_skipper = SchedulerRecvSkipper.maybe_create(self.server_args)
self.input_blocker = (
SchedulerInputBlocker(noop=self.attn_tp_rank != 0)
if get_bool_env_var("SGLANG_ENABLE_COLOCATED_BATCH_GEN")
else None
)
# Configure GC logger
if envs.SGLANG_LOG_GC.get():
configure_gc_logger()
def init_disaggregation(self):
self.disaggregation_mode = DisaggregationMode(
self.server_args.disaggregation_mode
@@ -966,6 +901,18 @@ class Scheduler(
# The prefill requests that are in the middle of kv sending
self.disagg_prefill_inflight_queue: List[Req] = []
# Init mm receiver for EPD disaggregation mode
if (
self.server_args.language_only
and self.server_args.encoder_transfer_backend == "zmq_to_scheduler"
):
self.mm_receiver = MMReceiver(
self.server_args,
hf_config=self.model_config.hf_config,
tp_rank=self.tp_rank,
pp_rank=self.pp_rank,
)
def init_overlap(self):
self.device_module = torch.get_device_module(self.device)
self.default_stream: CudaStream = self.device_module.current_stream()
@@ -995,15 +942,70 @@ class Scheduler(
self.batch_record_buf = [None] * 2
self.batch_record_ct = 0
def init_moe_config(self):
if hasattr(self.model_config.hf_config, "num_experts_per_tok"):
initialize_moe_config(self.server_args)
def init_deterministic_inference_config(self):
"""Initialize deterministic inference configuration for different attention backends."""
if not self.server_args.enable_deterministic_inference:
self.truncation_align_size = None
return
def init_gemm_config(self):
# Initialize GEMM-related configuration (currently FP8 Blockwise GEMM backend).
# Other GEMM backends (e.g. FP4, BF16, etc.) can be added here in the future.
# This is needed for FP8 quantization.
initialize_fp8_gemm_config(self.server_args)
backend_sizes = {
"flashinfer": ("SGLANG_FLASHINFER_PREFILL_SPLIT_TILE_SIZE", 4096),
"triton": ("SGLANG_TRITON_PREFILL_TRUNCATION_ALIGN_SIZE", 4096),
}
env_var, default_size = backend_sizes.get(
self.server_args.attention_backend, (None, None)
)
self.truncation_align_size = (
get_int_env_var(env_var, default_size) if env_var else None
)
def init_request_dispatcher(self):
self._request_dispatcher = TypeBasedDispatcher(
[
(TokenizedGenerateReqInput, self.handle_generate_request),
(TokenizedEmbeddingReqInput, self.handle_embedding_request),
(BatchTokenizedGenerateReqInput, self.handle_batch_generate_request),
(BatchTokenizedEmbeddingReqInput, self.handle_batch_embedding_request),
(FlushCacheReqInput, self.flush_cache_wrapped),
(ClearHiCacheReqInput, self.clear_hicache_storage_wrapped),
(AbortReq, self.abort_request),
(OpenSessionReqInput, self.open_session),
(CloseSessionReqInput, self.close_session),
(UpdateWeightFromDiskReqInput, self.update_weights_from_disk),
(InitWeightsUpdateGroupReqInput, self.init_weights_update_group),
(DestroyWeightsUpdateGroupReqInput, self.destroy_weights_update_group),
(
InitWeightsSendGroupForRemoteInstanceReqInput,
self.init_weights_send_group_for_remote_instance,
),
(
SendWeightsToRemoteInstanceReqInput,
self.send_weights_to_remote_instance,
),
(
UpdateWeightsFromDistributedReqInput,
self.update_weights_from_distributed,
),
(UpdateWeightsFromTensorReqInput, self.update_weights_from_tensor),
(UpdateWeightsFromIPCReqInput, self.update_weights_from_ipc),
(GetWeightsByNameReqInput, self.get_weights_by_name),
(ReleaseMemoryOccupationReqInput, self.release_memory_occupation),
(ResumeMemoryOccupationReqInput, self.resume_memory_occupation),
(CheckWeightsReqInput, self.check_weights),
(SlowDownReqInput, self.slow_down),
(ProfileReq, self.profile),
(FreezeGCReq, self.handle_freeze_gc),
(GetInternalStateReq, self.get_internal_state),
(SetInternalStateReq, self.set_internal_state),
(RpcReqInput, self.handle_rpc_request),
(ExpertDistributionReq, self.expert_distribution_handle),
(LoadLoRAAdapterReqInput, self.load_lora_adapter),
(UnloadLoRAAdapterReqInput, self.unload_lora_adapter),
(GetLoadReqInput, self.get_load),
(PauseGenerationReqInput, self.pause_generation),
(ContinueGenerationReqInput, self.continue_generation),
]
)
@DynamicGradMode()
def event_loop_normal(self):
@@ -1027,7 +1029,7 @@ class Scheduler(
# When the server is idle, do self-check and re-init some states
self.self_check_during_idle()
# Update the last batch
# Update last_batch
self.last_batch = batch
if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get():
self.self_check_during_busy()
@@ -1060,10 +1062,11 @@ class Scheduler(
pop_and_process()
# Launch the current batch
batch_result = None
if batch:
batch_result = self.run_batch(batch)
self.result_queue.append((batch.copy(), batch_result))
else:
batch_result = None
# Process the last batch
if self.last_batch:
@@ -1077,7 +1080,7 @@ class Scheduler(
# It depends on the result of the last batch (e.g., grammar), so we run it after the last batch is processed.
self.launch_batch_sample_if_needed(batch_result)
# Update the last batch
# Update last_batch
self.last_batch = batch
if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get():
self.self_check_during_busy()
@@ -1713,6 +1716,7 @@ class Scheduler(
# only finished requests to running_batch.
chunked_req_to_exclude.add(self.chunked_req)
self.tree_cache.cache_unfinished_req(self.chunked_req, chunked=True)
# chunked request keeps its rid but will get a new req_pool_idx
if self.tp_worker.model_runner.mambaish_config is not None:
self.req_to_token_pool.free(
@@ -1720,6 +1724,7 @@ class Scheduler(
)
else:
self.req_to_token_pool.free(self.chunked_req.req_pool_idx)
if self.last_batch and self.last_batch.forward_mode.is_extend():
if self.last_batch.chunked_req is not None:
# In the context pipeline parallelism, after the last chunk, the current microbatch still track outdated chunked_req.
@@ -2700,6 +2705,29 @@ def is_work_request(recv_req):
)
class SenderWrapper:
def __init__(self, socket: zmq.Socket):
self.socket = socket
def send_output(
self,
output: Union[BaseReq, BaseBatchReq],
recv_obj: Optional[Union[BaseReq, BaseBatchReq]] = None,
):
if self.socket is None:
return
if (
isinstance(recv_obj, BaseReq)
and recv_obj.http_worker_ipc is not None
and output.http_worker_ipc is None
):
# handle communicator reqs for multi-http worker case
output.http_worker_ipc = recv_obj.http_worker_ipc
self.socket.send_pyobj(output)
def run_scheduler_process(
server_args: ServerArgs,
port_args: PortArgs,
python/sglang/srt/managers/scheduler_metrics_mixin.py
+3
View File
@@ -80,6 +80,9 @@ class SchedulerMetricsMixin:
labels["dp_rank"] = dp_rank
self.metrics_collector = SchedulerMetricsCollector(labels=labels)
if self.enable_kv_cache_events:
self.init_kv_events(self.server_args.kv_events_config)
def init_kv_events(self: Scheduler, kv_events_config: Optional[str]):
if self.enable_kv_cache_events:
self.kv_event_publisher = EventPublisherFactory.create(