diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 83aeed081..b0f090530 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -823,20 +823,25 @@ class SchedulerDisaggregationDecodeMixin: """A normal scheduler loop for decode worker in disaggregation mode.""" while True: + # Receive requests recv_reqs = self.recv_requests() self.process_input_requests(recv_reqs) # polling and allocating kv cache self.process_decode_queue() + + # Get the next batch to run batch = self.get_next_disagg_decode_batch_to_run() self.cur_batch = batch + # Launch the current batch if batch: - # Generate fake extend output. result = self.run_batch(batch) self.process_batch_result(batch, result) else: + # When the server is idle, do self-check and re-init some states self.self_check_during_idle() + # Update last_batch self.last_batch = batch @torch.no_grad() @@ -845,26 +850,35 @@ class SchedulerDisaggregationDecodeMixin: self.last_batch: Optional[ScheduleBatch] = None while True: - + # Receive requests recv_reqs = self.recv_requests() self.process_input_requests(recv_reqs) # polling and allocating kv cache self.process_decode_queue() + + # Get the next batch to run batch = self.get_next_disagg_decode_batch_to_run() self.cur_batch = batch - batch_result = None + # Launch the current batch 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: tmp_batch, tmp_result = self.result_queue.popleft() self.process_batch_result(tmp_batch, tmp_result) elif batch is None: self.self_check_during_idle() + # Run sample of the current batch + # 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 last_batch self.last_batch = batch def _run_batch_prebuilt( diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index a1e04dbb7..bd0896b09 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -311,6 +311,10 @@ class SchedulerDisaggregationPrefillMixin: def get_next_disagg_prefill_batch_to_run( self: Scheduler, ) -> Optional[ScheduleBatch]: + # HACK (byronhsu): reset the batch_is_full flag because we never enter update_running_batch which resets it + # Otherwise, it hangs under high concurrency + self.running_batch.batch_is_full = False + self.process_prefill_chunk() batch = self.get_new_batch_prefill() @@ -327,14 +331,18 @@ class SchedulerDisaggregationPrefillMixin: """A normal scheduler loop for prefill worker in disaggregation mode.""" while True: + # Receive requests recv_reqs = self.recv_requests() self.process_input_requests(recv_reqs) self.waiting_queue.extend( self.disagg_prefill_bootstrap_queue.pop_bootstrapped() ) + + # Get the next batch to run batch = self.get_next_disagg_prefill_batch_to_run() self.cur_batch = batch + # Launch the current batch if batch: result = self.run_batch(batch) self.process_batch_result_disagg_prefill(batch, result) @@ -343,43 +351,48 @@ class SchedulerDisaggregationPrefillMixin: self.process_disagg_prefill_inflight_queue() + # Update last_batch self.last_batch = batch - # HACK (byronhsu): reset the batch_is_full flag because we never enter update_running_batch which resets it - # Otherwise, it hangs under high concurrency - self.running_batch.batch_is_full = False @torch.no_grad() def event_loop_overlap_disagg_prefill(self: Scheduler) -> None: self.result_queue = deque() while True: + # Receive requests recv_reqs = self.recv_requests() self.process_input_requests(recv_reqs) self.waiting_queue.extend( self.disagg_prefill_bootstrap_queue.pop_bootstrapped() ) + + # Get the next batch to run batch = self.get_next_disagg_prefill_batch_to_run() self.cur_batch = batch - batch_result = None + # Launch the current batch 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: tmp_batch, tmp_result = self.result_queue.popleft() self.process_batch_result_disagg_prefill(tmp_batch, tmp_result) elif batch is None: + # When the server is idle, do self-check and re-init some states self.self_check_during_idle() self.process_disagg_prefill_inflight_queue() + # Run sample of the current batch + # 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 last_batch self.last_batch = batch - # HACK (byronhsu): reset the batch_is_full flag because we never enter update_running_batch which resets it - # Otherwise, it hangs under high concurrency - self.running_batch.batch_is_full = False def process_batch_result_disagg_prefill( self: Scheduler, @@ -420,7 +433,6 @@ class SchedulerDisaggregationPrefillMixin: logits_output.input_token_logprobs.tolist() ) - hidden_state_offset = 0 for i, (req, next_token_id) in enumerate( zip(batch.reqs, next_token_ids, strict=True) ): diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index fa2d68240..354aa6c1d 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -91,7 +91,7 @@ class OpenAIServingChat(OpenAIServingBase): self.tokenizer_manager.tokenizer is not None and self.tokenizer_manager.tokenizer.chat_template is not None ) - architectures = self.tokenizer_manager.server_args.get_hf_config().architectures + architectures = self.tokenizer_manager.model_config.hf_config.architectures is_dpsk_v32 = "DeepseekV3" in architectures[0] if architectures else False return not has_chat_template and is_dpsk_v32 diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py b/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py index 9bb34046d..34494f599 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py @@ -1,7 +1,7 @@ import logging from sglang.srt.environ import envs -from sglang.srt.utils import get_device_sm, is_blackwell +from sglang.srt.utils import get_device_sm, is_blackwell_supported logger = logging.getLogger(__name__) @@ -21,5 +21,5 @@ def _compute_enable_deep_gemm(): ENABLE_JIT_DEEPGEMM = _compute_enable_deep_gemm() -DEEPGEMM_BLACKWELL = ENABLE_JIT_DEEPGEMM and is_blackwell() +DEEPGEMM_BLACKWELL = ENABLE_JIT_DEEPGEMM and is_blackwell_supported() DEEPGEMM_SCALE_UE8M0 = DEEPGEMM_BLACKWELL diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 45fbf5d65..033c37e3f 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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, diff --git a/python/sglang/srt/managers/scheduler_metrics_mixin.py b/python/sglang/srt/managers/scheduler_metrics_mixin.py index 59075c6ae..13a6c7159 100644 --- a/python/sglang/srt/managers/scheduler_metrics_mixin.py +++ b/python/sglang/srt/managers/scheduler_metrics_mixin.py @@ -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( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 3c9cd983d..27fc37880 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -26,8 +26,6 @@ import random import tempfile from typing import Any, Callable, Dict, List, Literal, Optional, Union -import orjson - from sglang.srt.connector import ConnectorType from sglang.srt.environ import ToolStrictLevel, envs from sglang.srt.function_call.function_call_parser import FunctionCallParser @@ -65,7 +63,7 @@ from sglang.srt.utils.common import ( wait_port_available, xpu_has_xmx_support, ) -from sglang.srt.utils.hf_transformers_utils import check_gguf_file, get_config +from sglang.srt.utils.hf_transformers_utils import check_gguf_file from sglang.utils import is_in_ci logger = logging.getLogger(__name__) @@ -190,7 +188,7 @@ FP8_GEMM_RUNNER_BACKEND_CHOICES = [ MAMBA_SSM_DTYPE_CHOICES = ["float32", "bfloat16"] -mamba_scheduler_strategy_CHOICES = ["auto", "no_buffer", "extra_buffer"] +MAMBA_SCHEDULER_STRATEGY_CHOICES = ["auto", "no_buffer", "extra_buffer"] # Allow external code to add more choices @@ -278,12 +276,6 @@ class ServerArgs: nccl_port: Optional[int] = None checkpoint_engine_wait_weights_before_ready: bool = False - # Encode prefill disaggregation - encoder_only: bool = False - language_only: bool = False - encoder_transfer_backend: str = ENCODER_TRANSFER_BACKEND_CHOICES[0] - encoder_urls: List[str] = dataclasses.field(default_factory=list) - # Quantization and data type dtype: str = "auto" quantization: Optional[str] = None @@ -609,6 +601,12 @@ class ServerArgs: # FIXME: hack to reduce ITL when decode bs is small disaggregation_decode_polling_interval: int = 1 + # Encode prefill disaggregation + encoder_only: bool = False + language_only: bool = False + encoder_transfer_backend: str = ENCODER_TRANSFER_BACKEND_CHOICES[0] + encoder_urls: List[str] = dataclasses.field(default_factory=list) + # For model weight update and weight loading custom_weight_loader: Optional[List[str]] = None weight_loader_disable_mmap: bool = False @@ -686,6 +684,7 @@ class ServerArgs: self._handle_a2a_moe() self._handle_eplb_and_dispatch() self._handle_expert_distribution_metrics() + self._handle_elastic_ep() # Handle pipeline parallelism. self._handle_pipeline_parallelism() @@ -693,9 +692,6 @@ class ServerArgs: # Handle speculative decoding logic. self._handle_speculative_decoding() - # Handle remote instance weight loader. - self._handle_remote_instance_weight_loader_start_seed_via_transfer_engine() - # Handle model loading format. self._handle_load_format() @@ -714,24 +710,15 @@ class ServerArgs: # Validate cache settings. self._handle_cache_compatibility() - # Validate metrics labels. - self._handle_metrics_labels() - # Handle deterministic inference. self._handle_deterministic_inference() - # Handle exporting request-level metrics. - self._handle_request_metrics_exporters() - # Handle diffusion LLM inference. self._handle_dllm_inference() # Handle any other necessary validations. self._handle_other_validations() - # Handle elastic expert parallelism. - self._handle_elastic_ep() - def _handle_deprecated_args(self): # Handle deprecated tool call parsers deprecated_tool_call_parsers = {"qwen25": "qwen", "glm45": "glm"} @@ -752,19 +739,8 @@ class ServerArgs: self.random_seed = random.randint(0, 1 << 30) if self.mm_process_config is None: self.mm_process_config = {} - if self.mamba_scheduler_strategy == "auto": - # TODO: when extra_buffer is more verified, we can set the default path based on - # [overlap, non-overlap] - self.mamba_scheduler_strategy = "no_buffer" - # In speculative scenario: - # - If `speculative_draft_model_quantization` is specified, the draft model uses this quantization method. - # - Otherwise, the draft model defaults to the same quantization as the target model. - if self.speculative_draft_model_quantization is None: - self.speculative_draft_model_quantization = self.quantization - elif self.speculative_draft_model_quantization == "unquant": - self.speculative_draft_model_quantization = None - # Handle ModelScope model downloads + # Handle ModelScope model downloads if get_bool_env_var("SGLANG_USE_MODELSCOPE"): if not os.path.exists(self.model_path): from modelscope import snapshot_download @@ -774,6 +750,44 @@ class ServerArgs: self.tokenizer_path, ignore_patterns=["*.bin", "*.safetensors"] ) + # Mamba scheduler strategy + if self.mamba_scheduler_strategy == "auto": + # TODO: when extra_buffer is more verified, we can set the default path based on + # [overlap, non-overlap] + self.mamba_scheduler_strategy = "no_buffer" + + # In speculative scenario: + # - If `speculative_draft_model_quantization` is specified, the draft model uses this quantization method. + # - Otherwise, the draft model defaults to the same quantization as the target model. + if self.speculative_draft_model_quantization is None: + self.speculative_draft_model_quantization = self.quantization + elif self.speculative_draft_model_quantization == "unquant": + self.speculative_draft_model_quantization = None + + def _handle_hpu_backends(self): + if self.device == "hpu": + self.attention_backend = "torch_native" + self.sampling_backend = "pytorch" + + def _handle_cpu_backends(self): + if self.device == "cpu": + if self.attention_backend is None: + self.attention_backend = "intel_amx" + self.sampling_backend = "pytorch" + + def _handle_npu_backends(self): + if self.device == "npu": + from sglang.srt.hardware_backend.npu.utils import set_default_server_args + + set_default_server_args(self) + + if self.piecewise_cuda_graph_compiler != "eager": + logger.warning( + "At this moment Ascend platform only support prefill graph compilation with " + "piecewise_cuda_graph_compiler='eager', change piecewise_cuda_graph_compiler to 'eager'." + ) + self.piecewise_cuda_graph_compiler = "eager" + def _handle_gpu_memory_settings(self, gpu_mem): """ Configure GPU memory-dependent settings including @@ -973,37 +987,13 @@ class ServerArgs: return capture_sizes - def _handle_hpu_backends(self): - if self.device == "hpu": - self.attention_backend = "torch_native" - self.sampling_backend = "pytorch" - - def _handle_cpu_backends(self): - if self.device == "cpu": - if self.attention_backend is None: - self.attention_backend = "intel_amx" - self.sampling_backend = "pytorch" - - def _handle_npu_backends(self): - if self.device == "npu": - from sglang.srt.hardware_backend.npu.utils import set_default_server_args - - set_default_server_args(self) - - if self.piecewise_cuda_graph_compiler != "eager": - logger.warning( - "At this moment Ascend platform only support prefill graph compilation with " - "piecewise_cuda_graph_compiler='eager', change piecewise_cuda_graph_compiler to 'eager'." - ) - self.piecewise_cuda_graph_compiler = "eager" - def _handle_model_specific_adjustments(self): from sglang.srt.configs.model_config import is_deepseek_nsa if parse_connector_type(self.model_path) == ConnectorType.INSTANCE: return - hf_config = self.get_hf_config() + hf_config = self.get_model_config().hf_config model_arch = hf_config.architectures[0] if model_arch in [ @@ -1017,16 +1007,13 @@ class ServerArgs: "MistralLarge3ForCausalLM", "PixtralForConditionalGeneration", ]: - if is_deepseek_nsa(hf_config): - if ( - self.attention_backend is None - and self.prefill_attention_backend is None - and self.decode_attention_backend is None - ): + # Set attention backend for DeepSeek + if is_deepseek_nsa(hf_config): # DeepSeek 3.2 + if self.is_attention_backend_not_set(): self.attention_backend = "nsa" - logger.warning("Set nsa attention backend for DeepSeek NSA.") + logger.info("Use nsa attention backend for DeepSeek NSA.") - if not is_npu(): + if not is_npu(): # CUDA GPU self.enable_dp_attention = True logger.warning("DP attention is enabled for DeepSeek NSA.") if self.enable_nsa_prefill_context_parallel: @@ -1090,10 +1077,11 @@ class ServerArgs: print_nsa_bool_env_vars() else: + # DeepSeek V3/R1/V3.1 if self.enable_piecewise_cuda_graph: logger.info("Piecewise CUDA graph is enabled, use MLA for prefill.") - if is_cuda() and is_sm100_supported(): + if is_sm100_supported(): if ( self.attention_backend is None and self.prefill_attention_backend is None @@ -1104,8 +1092,8 @@ class ServerArgs: "Use trtllm_mla as attention backend on sm100 for DeepseekV3ForCausalLM" ) - # common to all Deepseek MoE models - if is_cuda() and is_sm100_supported(): + # Set moe backend for DeepSeek + if is_sm100_supported(): quantization_config = getattr(hf_config, "quantization_config", None) quant_method = ( quantization_config.get("quant_method") @@ -1116,7 +1104,7 @@ class ServerArgs: # Default DeepSeek V3/R1 native FP8 when not explicitly set, # Because we need this condition for an assertion in # flashinfer_trtllm MoE runner backend. - if quant_method is None and model_arch == "DeepseekV3ForCausalLM": + if quant_method is None and model_arch in ["DeepseekV3ForCausalLM"]: self.quantization = "fp8" logger.info( "Quantization not specified, default to fp8 for DeepSeek on sm100" @@ -1134,14 +1122,11 @@ class ServerArgs: ) elif model_arch in ["GptOssForCausalLM"]: - if ( - self.attention_backend is None - and self.prefill_attention_backend is None - and self.decode_attention_backend is None - ): - if is_cuda() and is_sm100_supported(): + # Set attention backend for GPT-OSS + if self.is_attention_backend_not_set(): + if is_sm100_supported(): self.attention_backend = "trtllm_mha" - elif is_cuda() and is_sm90_supported(): + elif is_sm90_supported(): self.attention_backend = "fa3" else: self.attention_backend = "triton" @@ -1381,11 +1366,6 @@ class ServerArgs: FLA_CHUNK_SIZE % self.page_size == 0 ), f"Page size for hybrid GDN model must be divisible by {FLA_CHUNK_SIZE}, got {self.page_size}" - if self.speculative_algorithm is not None: - logger.info( - f"Disable overlap schedule for {model_arch} model speculative decoding." - ) - self.disable_overlap_schedule = True elif not self.disable_radix_cache: logger.warning( "Disabling overlap schedule since MambaRadixCache no_buffer is not compatible with " @@ -1421,6 +1401,10 @@ class ServerArgs: self.disable_radix_cache = True self.disable_overlap_schedule = False + if not self.get_model_config().is_generation: + self.disable_overlap_schedule = True + logger.warning("Overlap scheduler is disabled for embedding models.") + # TRTLLM AllReduce Fusion supports SM90/100/120, enable it by default # for models with explicit support (DeepseekV3, GptOss, Glm4Moe, Qwen3Moe) # TODO: currently, it is only supported in the single node scenario. https://github.com/flashinfer-ai/flashinfer/issues/2006 @@ -1445,9 +1429,6 @@ class ServerArgs: and self.moe_a2a_backend == "none" ): self.enable_flashinfer_allreduce_fusion = True - logger.info( - f"Enable FlashInfer AllReduce Fusion by default for {model_arch}" - ) def _handle_sampling_backend(self): if self.sampling_backend is None: @@ -1520,7 +1501,7 @@ class ServerArgs: self.attention_backend = "triton" logger.warning( - f"Attention backend not explicitly specified. Use {self.attention_backend} backend by default." + f"Attention backend not specified. Use {self.attention_backend} backend by default." ) # Torch native and flex attention backends @@ -1977,7 +1958,7 @@ class ServerArgs: "eagle speculative decoding." ) - model_arch = self.get_hf_config().architectures[0] + model_arch = self.get_model_config().hf_config.architectures[0] if model_arch in [ "DeepseekV32ForCausalLM", "DeepseekV3ForCausalLM", @@ -2117,27 +2098,10 @@ class ServerArgs: ) self.load_format = "auto" - def _handle_encoder_disaggregation(self): - if self.enable_prefix_mm_cache and not self.encoder_only: - raise ValueError( - "--enable-prefix-mm-cache requires --encoder-only to be enabled" - ) - if self.encoder_only and self.language_only: - raise ValueError("Cannot set --encoder-only and --language-only together") - if self.encoder_only and not self.disaggregation_mode == "null": - raise ValueError( - "Cannot set --encoder-only and --disaggregation-mode prefill/decode together" - ) - if ( - self.language_only - and self.encoder_transfer_backend == "zmq_to_scheduler" - and self.pp_size > 1 - ): - raise ValueError("zmq_to_scheduler not support pp_size > 1") - - if self.language_only and len(self.encoder_urls) == 0: - raise ValueError( - "requires at least one encoder urls to be set via --encoder-urls" + # Check whether TransferEngine can be used when users want to start seed service that supports TransferEngine backend. + if self.remote_instance_weight_loader_start_seed_via_transfer_engine: + self.remote_instance_weight_loader_start_seed_via_transfer_engine = ( + self.validate_transfer_engine() ) def _handle_pd_disaggregation(self): @@ -2173,6 +2137,29 @@ class ServerArgs: "Cuda graph is disabled for prefill server when piecewise cuda graph is not enabled." ) + def _handle_encoder_disaggregation(self): + if self.enable_prefix_mm_cache and not self.encoder_only: + raise ValueError( + "--enable-prefix-mm-cache requires --encoder-only to be enabled" + ) + if self.encoder_only and self.language_only: + raise ValueError("Cannot set --encoder-only and --language-only together") + if self.encoder_only and not self.disaggregation_mode == "null": + raise ValueError( + "Cannot set --encoder-only and --disaggregation-mode prefill/decode together" + ) + if ( + self.language_only + and self.encoder_transfer_backend == "zmq_to_scheduler" + and self.pp_size > 1 + ): + raise ValueError("zmq_to_scheduler not support pp_size > 1") + + if self.language_only and len(self.encoder_urls) == 0: + raise ValueError( + "requires at least one encoder urls to be set via --encoder-urls" + ) + def _handle_tokenizer_batching(self): if self.enable_tokenizer_batch_encode and self.enable_dynamic_batch_tokenizer: raise ValueError( @@ -2238,15 +2225,6 @@ class ServerArgs: "Spec v2 and decode offload kv cache are incompatible and cannot be enabled together." ) - def _handle_metrics_labels(self): - if ( - not self.tokenizer_metrics_custom_labels_header - and self.tokenizer_metrics_allowed_custom_labels - ): - raise ValueError( - "Please set --tokenizer-metrics-custom-labels-header when setting --tokenizer-metrics-allowed-custom-labels." - ) - def _handle_deterministic_inference(self): if self.rl_on_policy_target is not None: logger.warning( @@ -2265,7 +2243,7 @@ class ServerArgs: is_deepseek_model = False if parse_connector_type(self.model_path) != ConnectorType.INSTANCE: try: - hf_config = self.get_hf_config() + hf_config = self.get_model_config().hf_config model_arch = hf_config.architectures[0] is_deepseek_model = model_arch in [ "DeepseekV2ForCausalLM", @@ -2326,13 +2304,6 @@ class ServerArgs: "NCCL_ALGO is set to 'allreduce:tree' and custom all reduce is disabled for deterministic inference when TP size > 1." ) - def _handle_request_metrics_exporters(self): - """Handle arguments for configuring `RequestMetricsExporter` usage.""" - if self.export_metrics_to_file and self.export_metrics_to_file_dir is None: - raise ValueError( - "--export-metrics-to-file-dir is required when --export-metrics-to-file is enabled" - ) - def _handle_dllm_inference(self): if self.dllm_algorithm is None: return @@ -2372,13 +2343,6 @@ class ServerArgs: self.disable_cuda_graph = True self.skip_server_warmup = True - def _handle_remote_instance_weight_loader_start_seed_via_transfer_engine(self): - # Check whether TransferEngine can be used when users want to start seed service that supports TransferEngine backend. - if self.remote_instance_weight_loader_start_seed_via_transfer_engine: - self.remote_instance_weight_loader_start_seed_via_transfer_engine = ( - self.validate_transfer_engine() - ) - @staticmethod def add_cli_args(parser: argparse.ArgumentParser): @@ -2538,32 +2502,6 @@ class ServerArgs: "before serving inference requests.", ) - # Encode prefill disaggregation - parser.add_argument( - "--encoder-only", - action="store_true", - help="For MLLM with an encoder, launch an encoder-only server", - ) - parser.add_argument( - "--language-only", - action="store_true", - help="For VLM, load weights for the language model only.", - ) - parser.add_argument( - "--encoder-transfer-backend", - type=str, - default=ServerArgs.encoder_transfer_backend, - choices=ENCODER_TRANSFER_BACKEND_CHOICES, - help="The backend for encoder disaggregation transfer. Default is zmq_to_scheduler.", - ) - parser.add_argument( - "--encoder-urls", - nargs="+", - type=str, - default=[], - help="List of encoder server urls.", - ) - # Quantization and data type parser.add_argument( "--dtype", @@ -3617,7 +3555,7 @@ class ServerArgs: parser.add_argument( "--mamba-scheduler-strategy", type=str, - choices=mamba_scheduler_strategy_CHOICES, + choices=MAMBA_SCHEDULER_STRATEGY_CHOICES, default=ServerArgs.mamba_scheduler_strategy, help="The strategy to use for mamba radix cache.", ) @@ -4245,6 +4183,32 @@ class ServerArgs: help="The interval to poll requests in decode server. Can be set to >1 to reduce the overhead of this.", ) + # Encode prefill disaggregation + parser.add_argument( + "--encoder-only", + action="store_true", + help="For MLLM with an encoder, launch an encoder-only server", + ) + parser.add_argument( + "--language-only", + action="store_true", + help="For VLM, load weights for the language model only.", + ) + parser.add_argument( + "--encoder-transfer-backend", + type=str, + default=ServerArgs.encoder_transfer_backend, + choices=ENCODER_TRANSFER_BACKEND_CHOICES, + help="The backend for encoder disaggregation transfer. Default is zmq_to_scheduler.", + ) + parser.add_argument( + "--encoder-urls", + nargs="+", + type=str, + default=[], + help="List of encoder server urls.", + ) + # Custom weight loader parser.add_argument( "--custom-weight-loader", @@ -4391,17 +4355,6 @@ class ServerArgs: else: return f"http://{self.host}:{self.port}" - def get_hf_config(self): - kwargs = {} - hf_config = get_config( - self.model_path, - trust_remote_code=self.trust_remote_code, - revision=self.revision, - model_override_args=orjson.loads(self.json_model_override_args), - **kwargs, - ) - return hf_config - def get_model_config(self): # Lazy init to avoid circular import from sglang.srt.configs.model_config import ModelConfig @@ -4430,6 +4383,13 @@ class ServerArgs: model_config = self.get_model_config() return model_config.attention_arch == AttentionArch.MLA + def is_attention_backend_not_set(self): + return ( + self.attention_backend is None + and self.prefill_attention_backend is None + and self.decode_attention_backend is None + ) + def enable_mamba_extra_buffer(self) -> bool: return self.mamba_scheduler_strategy == "extra_buffer" @@ -4548,6 +4508,21 @@ class ServerArgs: if self.model_impl == "mindspore": assert is_npu(), "MindSpore model impl is only supported on Ascend npu." + # Check metrics labels + if ( + not self.tokenizer_metrics_custom_labels_header + and self.tokenizer_metrics_allowed_custom_labels + ): + raise ValueError( + "Please set --tokenizer-metrics-custom-labels-header when setting --tokenizer-metrics-allowed-custom-labels." + ) + + # Check metrics exporters + if self.export_metrics_to_file and self.export_metrics_to_file_dir is None: + raise ValueError( + "--export-metrics-to-file-dir is required when --export-metrics-to-file is enabled" + ) + def check_torch_2_9_1_cudnn_compatibility(self): if get_bool_env_var("SGLANG_DISABLE_CUDNN_CHECK"): return @@ -5035,7 +5010,7 @@ def auto_choose_speculative_params(self: ServerArgs): You can tune them on your own models and prompts with scripts/playground/bench_speculative.py """ - hf_config = self.get_hf_config() + hf_config = self.get_model_config().hf_config arch = hf_config.architectures[0] if self.speculative_algorithm == "STANDALONE": # The default value for standalone speculative decoding diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 2e1986ba8..b52e9483d 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -49,7 +49,7 @@ import warnings from collections import OrderedDict, defaultdict from contextlib import contextmanager from dataclasses import dataclass -from functools import lru_cache +from functools import lru_cache, partial from importlib.metadata import PackageNotFoundError, version from importlib.util import find_spec from io import BytesIO @@ -105,22 +105,6 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -show_time_cost = False -time_infos = {} - - -def get_or_create_event_loop(): - """Gets the running event loop or creates a new one if it doesn't exist.""" - try: - return asyncio.get_running_loop() - except RuntimeError: - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - return loop - - -HIP_FP8_E4M3_FNUZ_MAX = 224.0 - # https://pytorch.org/docs/stable/notes/hip.html#checking-for-hip @lru_cache(maxsize=1) @@ -129,6 +113,7 @@ def is_hip() -> bool: if is_hip(): + HIP_FP8_E4M3_FNUZ_MAX = 224.0 FP8_E4M3_MAX = HIP_FP8_E4M3_FNUZ_MAX else: FP8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max @@ -191,14 +176,6 @@ def get_cuda_version(): return (0, 0) -def _check(cc_major): - if not is_cuda(): - return False - return torch.cuda.get_device_capability()[0] == cc_major and tuple( - map(int, torch.version.cuda.split(".")[:2]) - ) >= (12, 3) - - @contextmanager def device_context(device: torch.device): if device.type == "cpu" and is_cpu(): @@ -213,96 +190,49 @@ def device_context(device: torch.device): raise ValueError(f"Unknown device module: {device}") -is_ampere_with_cuda_12_3 = lambda: _check(8) -is_hopper_with_cuda_12_3 = lambda: _check(9) - - -@lru_cache(maxsize=1) -def is_blackwell(): +def _check_cuda_device_version( + device_capability_majors: List[int], cuda_version: Tuple[int, int] +): if not is_cuda(): return False - return torch.cuda.get_device_capability()[0] in [10, 12] - - -@lru_cache(maxsize=1) -def is_blackwell_supported(device=None) -> bool: - if not is_cuda(): - return False - return is_sm100_supported(device) or is_sm120_supported(device) - - -@lru_cache(maxsize=1) -def is_sm120_supported(device=None) -> bool: - if not is_cuda(): - return False - return (torch.cuda.get_device_capability(device)[0] == 12) and ( - torch.version.cuda >= "12.8" + return ( + torch.cuda.get_device_capability()[0] in device_capability_majors + and tuple(map(int, torch.version.cuda.split(".")[:2])) >= cuda_version ) -@lru_cache(maxsize=1) -def is_sm100_supported(device=None) -> bool: - if not is_cuda(): - return False - return (torch.cuda.get_device_capability(device)[0] == 10) and ( - torch.version.cuda >= "12.8" +is_ampere_with_cuda_12_3 = lru_cache(maxsize=1)( + partial( + _check_cuda_device_version, device_capability_majors=[8], cuda_version=(12, 3) ) - - -@lru_cache(maxsize=1) -def is_sm90_supported(device=None) -> bool: - if not is_cuda(): - return False - return (torch.cuda.get_device_capability(device)[0] == 9) and ( - torch.version.cuda >= "12.3" +) +is_hopper_with_cuda_12_3 = lru_cache(maxsize=1)( + partial( + _check_cuda_device_version, device_capability_majors=[9], cuda_version=(12, 3) ) - - -_warned_bool_env_var_keys = set() - - -def get_bool_env_var(name: str, default: str = "false") -> bool: - # FIXME: move your environment variable to sglang.srt.environ - value = os.getenv(name, default) - value = value.lower() - - truthy_values = ("true", "1") - falsy_values = ("false", "0") - - if (value not in truthy_values) and (value not in falsy_values): - if value not in _warned_bool_env_var_keys: - logger.warning( - f"get_bool_env_var({name}) see non-understandable value={value} and treat as false" - ) - _warned_bool_env_var_keys.add(value) - - return value in truthy_values - - -def get_int_env_var(name: str, default: int = 0) -> int: - # FIXME: move your environment variable to sglang.srt.environ - value = os.getenv(name) - if value is None or not value.strip(): - return default - try: - return int(value) - except ValueError: - return default - - -def get_float_env_var(name: str, default: float = 0.0) -> float: - # FIXME: move your environment variable to sglang.srt.environ - value = os.getenv(name) - if value is None or not value.strip(): - return default - try: - return float(value) - except ValueError: - return default - - -def support_triton(backend: str) -> bool: - return backend not in ["torch_native", "intel_amx"] +) +is_blackwell_supported = is_blackwell = lru_cache(maxsize=1)( + partial( + _check_cuda_device_version, + device_capability_majors=[10, 12], + cuda_version=(12, 8), + ) +) +is_sm120_supported = lru_cache(maxsize=1)( + partial( + _check_cuda_device_version, device_capability_majors=[12], cuda_version=(12, 8) + ) +) +is_sm100_supported = lru_cache(maxsize=1)( + partial( + _check_cuda_device_version, device_capability_majors=[10], cuda_version=(12, 8) + ) +) +is_sm90_supported = lru_cache(maxsize=1)( + partial( + _check_cuda_device_version, device_capability_majors=[9], cuda_version=(12, 3) + ) +) try: @@ -364,6 +294,53 @@ def random_uuid() -> str: return str(uuid.uuid4().hex) +_warned_bool_env_var_keys = set() + + +def get_bool_env_var(name: str, default: str = "false") -> bool: + # FIXME: move your environment variable to sglang.srt.environ + value = os.getenv(name, default) + value = value.lower() + + truthy_values = ("true", "1") + falsy_values = ("false", "0") + + if (value not in truthy_values) and (value not in falsy_values): + if value not in _warned_bool_env_var_keys: + logger.warning( + f"get_bool_env_var({name}) see non-understandable value={value} and treat as false" + ) + _warned_bool_env_var_keys.add(value) + + return value in truthy_values + + +def get_int_env_var(name: str, default: int = 0) -> int: + # FIXME: move your environment variable to sglang.srt.environ + value = os.getenv(name) + if value is None or not value.strip(): + return default + try: + return int(value) + except ValueError: + return default + + +def get_float_env_var(name: str, default: float = 0.0) -> float: + # FIXME: move your environment variable to sglang.srt.environ + value = os.getenv(name) + if value is None or not value.strip(): + return default + try: + return float(value) + except ValueError: + return default + + +def support_triton(backend: str) -> bool: + return backend not in ["torch_native", "intel_amx"] + + _ENABLE_TORCH_INFERENCE_MODE = get_bool_env_var( "SGLANG_ENABLE_TORCH_INFERENCE_MODE", "false" ) @@ -421,6 +398,10 @@ class DynamicGradMode(_DecoratorContextManager): return self.__class__() +show_time_cost = False +time_infos = {} + + def enable_show_time_cost(): global show_time_cost show_time_cost = True @@ -3778,3 +3759,13 @@ def raise_error_or_warn(obj, strict, counter_name, message, log_interval=1000): if count % log_interval == 0: logger.warning(message) setattr(obj, counter_name, count + 1) + + +def get_or_create_event_loop(): + """Gets the running event loop or creates a new one if it doesn't exist.""" + try: + return asyncio.get_running_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + return loop diff --git a/test/srt/openai_server/basic/test_serving_chat.py b/test/srt/openai_server/basic/test_serving_chat.py index ed20cbc94..20acc8605 100644 --- a/test/srt/openai_server/basic/test_serving_chat.py +++ b/test/srt/openai_server/basic/test_serving_chat.py @@ -36,7 +36,7 @@ class _MockTokenizerManager: # Mock hf_config for _use_dpsk_v32_encoding check mock_hf_config = Mock() mock_hf_config.architectures = ["LlamaForCausalLM"] - self.server_args.get_hf_config.return_value = mock_hf_config + self.model_config.hf_config = mock_hf_config self.chat_template_name: Optional[str] = "llama-3" @@ -603,9 +603,7 @@ class ServingChatTestCase(unittest.TestCase): mock_hf_config = Mock() mock_hf_config.architectures = ["DeepseekV32ForCausalLM"] - tokenizer_manager.server_args.get_hf_config = Mock( - return_value=mock_hf_config - ) + tokenizer_manager.model_config.hf_config = mock_hf_config # Case 1: No chat template in tokenizer -> should use dpsk encoding tokenizer_manager.tokenizer = Mock()