chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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/disaggregation/decode.py
+17 -3
View File
@@ -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(
python/sglang/srt/disaggregation/prefill.py
+20 -8
View File
@@ -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)
):
python/sglang/srt/entrypoints/openai/serving_chat.py
+1 -1
View File
@@ -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
python/sglang/srt/layers/deep_gemm_wrapper/configurer.py
+2 -2
View File
@@ -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
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(
python/sglang/srt/server_args.py
+147 -172
View File
@@ -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
python/sglang/srt/utils/common.py
+99 -108
View File
@@ -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