minke.yu/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.""" """A normal scheduler loop for decode worker in disaggregation mode."""
while True: while True:
# Receive requests
recv_reqs = self.recv_requests() recv_reqs = self.recv_requests()
self.process_input_requests(recv_reqs) self.process_input_requests(recv_reqs)
# polling and allocating kv cache # polling and allocating kv cache
self.process_decode_queue() self.process_decode_queue()
# Get the next batch to run
batch = self.get_next_disagg_decode_batch_to_run() batch = self.get_next_disagg_decode_batch_to_run()
self.cur_batch = batch self.cur_batch = batch
# Launch the current batch
if batch: if batch:
# Generate fake extend output.
result = self.run_batch(batch) result = self.run_batch(batch)
self.process_batch_result(batch, result) self.process_batch_result(batch, result)
else: else:
# When the server is idle, do self-check and re-init some states
self.self_check_during_idle() self.self_check_during_idle()
# Update last_batch
self.last_batch = batch self.last_batch = batch
@torch.no_grad() @torch.no_grad()
@@ -845,26 +850,35 @@ class SchedulerDisaggregationDecodeMixin:
self.last_batch: Optional[ScheduleBatch] = None self.last_batch: Optional[ScheduleBatch] = None
while True: while True:
# Receive requests
recv_reqs = self.recv_requests() recv_reqs = self.recv_requests()
self.process_input_requests(recv_reqs) self.process_input_requests(recv_reqs)
# polling and allocating kv cache # polling and allocating kv cache
self.process_decode_queue() self.process_decode_queue()
# Get the next batch to run
batch = self.get_next_disagg_decode_batch_to_run() batch = self.get_next_disagg_decode_batch_to_run()
self.cur_batch = batch self.cur_batch = batch
batch_result = None # Launch the current batch
if batch: if batch:
batch_result = self.run_batch(batch) batch_result = self.run_batch(batch)
self.result_queue.append((batch.copy(), batch_result)) self.result_queue.append((batch.copy(), batch_result))
else:
batch_result = None
# Process the last batch
if self.last_batch: if self.last_batch:
tmp_batch, tmp_result = self.result_queue.popleft() tmp_batch, tmp_result = self.result_queue.popleft()
self.process_batch_result(tmp_batch, tmp_result) self.process_batch_result(tmp_batch, tmp_result)
elif batch is None: elif batch is None:
self.self_check_during_idle() 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) self.launch_batch_sample_if_needed(batch_result)
# Update last_batch
self.last_batch = batch self.last_batch = batch
def _run_batch_prebuilt( 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( def get_next_disagg_prefill_batch_to_run(
self: Scheduler, self: Scheduler,
) -> Optional[ScheduleBatch]: ) -> 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() self.process_prefill_chunk()
batch = self.get_new_batch_prefill() batch = self.get_new_batch_prefill()
@@ -327,14 +331,18 @@ class SchedulerDisaggregationPrefillMixin:
"""A normal scheduler loop for prefill worker in disaggregation mode.""" """A normal scheduler loop for prefill worker in disaggregation mode."""
while True: while True:
# Receive requests
recv_reqs = self.recv_requests() recv_reqs = self.recv_requests()
self.process_input_requests(recv_reqs) self.process_input_requests(recv_reqs)
self.waiting_queue.extend( self.waiting_queue.extend(
self.disagg_prefill_bootstrap_queue.pop_bootstrapped() self.disagg_prefill_bootstrap_queue.pop_bootstrapped()
) )
# Get the next batch to run
batch = self.get_next_disagg_prefill_batch_to_run() batch = self.get_next_disagg_prefill_batch_to_run()
self.cur_batch = batch self.cur_batch = batch
# Launch the current batch
if batch: if batch:
result = self.run_batch(batch) result = self.run_batch(batch)
self.process_batch_result_disagg_prefill(batch, result) self.process_batch_result_disagg_prefill(batch, result)
@@ -343,43 +351,48 @@ class SchedulerDisaggregationPrefillMixin:
self.process_disagg_prefill_inflight_queue() self.process_disagg_prefill_inflight_queue()
# Update last_batch
self.last_batch = 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() @torch.no_grad()
def event_loop_overlap_disagg_prefill(self: Scheduler) -> None: def event_loop_overlap_disagg_prefill(self: Scheduler) -> None:
self.result_queue = deque() self.result_queue = deque()
while True: while True:
# Receive requests
recv_reqs = self.recv_requests() recv_reqs = self.recv_requests()
self.process_input_requests(recv_reqs) self.process_input_requests(recv_reqs)
self.waiting_queue.extend( self.waiting_queue.extend(
self.disagg_prefill_bootstrap_queue.pop_bootstrapped() self.disagg_prefill_bootstrap_queue.pop_bootstrapped()
) )
# Get the next batch to run
batch = self.get_next_disagg_prefill_batch_to_run() batch = self.get_next_disagg_prefill_batch_to_run()
self.cur_batch = batch self.cur_batch = batch
batch_result = None # Launch the current batch
if batch: if batch:
batch_result = self.run_batch(batch) batch_result = self.run_batch(batch)
self.result_queue.append((batch.copy(), batch_result)) self.result_queue.append((batch.copy(), batch_result))
else:
batch_result = None
# Process the last batch
if self.last_batch: if self.last_batch:
tmp_batch, tmp_result = self.result_queue.popleft() tmp_batch, tmp_result = self.result_queue.popleft()
self.process_batch_result_disagg_prefill(tmp_batch, tmp_result) self.process_batch_result_disagg_prefill(tmp_batch, tmp_result)
elif batch is None: elif batch is None:
# When the server is idle, do self-check and re-init some states
self.self_check_during_idle() self.self_check_during_idle()
self.process_disagg_prefill_inflight_queue() 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) self.launch_batch_sample_if_needed(batch_result)
# Update last_batch
self.last_batch = 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( def process_batch_result_disagg_prefill(
self: Scheduler, self: Scheduler,
@@ -420,7 +433,6 @@ class SchedulerDisaggregationPrefillMixin:
logits_output.input_token_logprobs.tolist() logits_output.input_token_logprobs.tolist()
) )
hidden_state_offset = 0
for i, (req, next_token_id) in enumerate( for i, (req, next_token_id) in enumerate(
zip(batch.reqs, next_token_ids, strict=True) 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 self.tokenizer_manager.tokenizer is not None
and self.tokenizer_manager.tokenizer.chat_template 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 is_dpsk_v32 = "DeepseekV3" in architectures[0] if architectures else False
return not has_chat_template and is_dpsk_v32 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 import logging
from sglang.srt.environ import envs 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__) logger = logging.getLogger(__name__)
@@ -21,5 +21,5 @@ def _compute_enable_deep_gemm():
ENABLE_JIT_DEEPGEMM = _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 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.moe_ep_size = server_args.ep_size
self.pp_size = server_args.pp_size self.pp_size = server_args.pp_size
self.dp_size = server_args.dp_size self.dp_size = server_args.dp_size
self.nccl_port = port_args.nccl_port
self.schedule_policy = server_args.schedule_policy self.schedule_policy = server_args.schedule_policy
self.enable_priority_scheduling = server_args.enable_priority_scheduling self.enable_priority_scheduling = server_args.enable_priority_scheduling
self.abort_on_priority_when_disabled = ( self.abort_on_priority_when_disabled = (
@@ -292,9 +293,11 @@ class Scheduler(
# Init model config # Init model config
self.model_config = ModelConfig.from_server_args(server_args) self.model_config = ModelConfig.from_server_args(server_args)
self.dllm_config = ( # For diffusion LLM
# Init diffusion LLM config DllmConfig.from_server_args(server_args)
self.dllm_config = DllmConfig.from_server_args(server_args) if server_args.dllm_algorithm is not None
else None
)
# Init metrics stats # Init metrics stats
self.init_metrics(tp_rank, pp_rank, dp_rank) self.init_metrics(tp_rank, pp_rank, dp_rank)
@@ -309,348 +312,49 @@ class Scheduler(
# Init tokenizer # Init tokenizer
self.init_tokenizer() self.init_tokenizer()
# Init moe config # Init moe config and GEMM config (FP8 GEMM, etc.)
self.init_moe_config() self.init_moe_gemm_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.")
# Launch a tensor parallel worker # Launch a tensor parallel worker
from sglang.srt.managers.tp_worker import TpModelWorker self.init_model_worker()
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"
)
# Init cache using the existing memory pool # Init cache using the existing memory pool
self.init_cache_with_memory_pool() self.init_cache_with_memory_pool()
# Init running status # Init running status
self.waiting_queue: List[Req] = [] self.init_running_status()
# 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
# Init chunked prefill # Init chunked prefill
self.chunked_prefill_size = server_args.chunked_prefill_size self.init_chunked_prefill()
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
# Init the grammar backend for constrained generation # Init the grammar backend for constrained generation
self.grammar_queue: List[Req] = [] self.init_grammar_backend()
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
# Init schedule policy and new token estimation # Init schedule policy and new token estimation
self.policy = SchedulePolicy( self.init_schedule_policy()
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
# Init watchdog thread # Init watchdog, memory saver, input blocker and recv skipper
self.watchdog = SchedulerWatchdog( self.init_watch_dog_memory_saver_input_blocker()
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 memory saver, profiler and metric stats # Init profiler
self.memory_saver_adapter = TorchMemorySaverAdapter.create(
enable=server_args.enable_memory_saver
)
self.offload_tags = set()
self.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 # Init disaggregation
self.init_disaggregation() self.init_disaggregation()
if self.enable_kv_cache_events: # Init overlap
self.init_kv_events(server_args.kv_events_config) self.init_overlap()
if envs.SGLANG_LOG_GC.get():
configure_gc_logger()
# Init prefill kv split size when deterministic inference is enabled with various attention backends # Init prefill kv split size when deterministic inference is enabled with various attention backends
self.init_deterministic_inference_config() 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 # Init request dispatcher
self._request_dispatcher = TypeBasedDispatcher( self.init_request_dispatcher()
[
(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),
]
)
def init_sockets(self, server_args: ServerArgs, port_args: PortArgs): def init_sockets(self, server_args: ServerArgs, port_args: PortArgs):
context = zmq.Context(2) context = zmq.Context(2)
self.idle_sleeper = None 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: if self.pp_rank == 0 and self.attn_tp_rank == 0:
self.recv_from_tokenizer = get_zmq_socket( self.recv_from_tokenizer = get_zmq_socket(
context, zmq.PULL, port_args.scheduler_input_ipc_name, False 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 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): def init_tokenizer(self):
server_args = self.server_args server_args = self.server_args
self.is_generation = self.model_config.is_generation self.is_generation = self.model_config.is_generation
@@ -744,13 +431,141 @@ class Scheduler(
reasoning_parser.detector.think_end_token, add_special_tokens=False reasoning_parser.detector.think_end_token, add_special_tokens=False
)[0] )[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): def init_cache_with_memory_pool(self):
server_args = self.server_args 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.req_to_token_pool, self.token_to_kv_pool_allocator = (
self.tp_worker.get_memory_pool() self.tp_worker.get_memory_pool()
) )
# Create cache
params = CacheInitParams( params = CacheInitParams(
disable=server_args.disable_radix_cache, disable=server_args.disable_radix_cache,
req_to_token_pool=self.req_to_token_pool, 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() embedding_cache_size = envs.SGLANG_VLM_CACHE_SIZE_MB.get()
init_mm_embedding_cache(embedding_cache_size * 1024 * 1024) 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): def init_disaggregation(self):
self.disaggregation_mode = DisaggregationMode( self.disaggregation_mode = DisaggregationMode(
self.server_args.disaggregation_mode self.server_args.disaggregation_mode
@@ -966,6 +901,18 @@ class Scheduler(
# The prefill requests that are in the middle of kv sending # The prefill requests that are in the middle of kv sending
self.disagg_prefill_inflight_queue: List[Req] = [] 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): def init_overlap(self):
self.device_module = torch.get_device_module(self.device) self.device_module = torch.get_device_module(self.device)
self.default_stream: CudaStream = self.device_module.current_stream() self.default_stream: CudaStream = self.device_module.current_stream()
@@ -995,15 +942,70 @@ class Scheduler(
self.batch_record_buf = [None] * 2 self.batch_record_buf = [None] * 2
self.batch_record_ct = 0 self.batch_record_ct = 0
def init_moe_config(self): def init_deterministic_inference_config(self):
if hasattr(self.model_config.hf_config, "num_experts_per_tok"): """Initialize deterministic inference configuration for different attention backends."""
initialize_moe_config(self.server_args) if not self.server_args.enable_deterministic_inference:
self.truncation_align_size = None
return
def init_gemm_config(self): backend_sizes = {
# Initialize GEMM-related configuration (currently FP8 Blockwise GEMM backend). "flashinfer": ("SGLANG_FLASHINFER_PREFILL_SPLIT_TILE_SIZE", 4096),
# Other GEMM backends (e.g. FP4, BF16, etc.) can be added here in the future. "triton": ("SGLANG_TRITON_PREFILL_TRUNCATION_ALIGN_SIZE", 4096),
# This is needed for FP8 quantization. }
initialize_fp8_gemm_config(self.server_args) 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() @DynamicGradMode()
def event_loop_normal(self): def event_loop_normal(self):
@@ -1027,7 +1029,7 @@ class Scheduler(
# When the server is idle, do self-check and re-init some states # When the server is idle, do self-check and re-init some states
self.self_check_during_idle() self.self_check_during_idle()
# Update the last batch # Update last_batch
self.last_batch = batch self.last_batch = batch
if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get(): if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get():
self.self_check_during_busy() self.self_check_during_busy()
@@ -1060,10 +1062,11 @@ class Scheduler(
pop_and_process() pop_and_process()
# Launch the current batch # Launch the current batch
batch_result = None
if batch: if batch:
batch_result = self.run_batch(batch) batch_result = self.run_batch(batch)
self.result_queue.append((batch.copy(), batch_result)) self.result_queue.append((batch.copy(), batch_result))
else:
batch_result = None
# Process the last batch # Process the last batch
if self.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. # 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) self.launch_batch_sample_if_needed(batch_result)
# Update the last batch # Update last_batch
self.last_batch = batch self.last_batch = batch
if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get(): if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get():
self.self_check_during_busy() self.self_check_during_busy()
@@ -1713,6 +1716,7 @@ class Scheduler(
# only finished requests to running_batch. # only finished requests to running_batch.
chunked_req_to_exclude.add(self.chunked_req) chunked_req_to_exclude.add(self.chunked_req)
self.tree_cache.cache_unfinished_req(self.chunked_req, chunked=True) self.tree_cache.cache_unfinished_req(self.chunked_req, chunked=True)
# chunked request keeps its rid but will get a new req_pool_idx # chunked request keeps its rid but will get a new req_pool_idx
if self.tp_worker.model_runner.mambaish_config is not None: if self.tp_worker.model_runner.mambaish_config is not None:
self.req_to_token_pool.free( self.req_to_token_pool.free(
@@ -1720,6 +1724,7 @@ class Scheduler(
) )
else: else:
self.req_to_token_pool.free(self.chunked_req.req_pool_idx) 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 and self.last_batch.forward_mode.is_extend():
if self.last_batch.chunked_req is not None: 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. # 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( def run_scheduler_process(
server_args: ServerArgs, server_args: ServerArgs,
port_args: PortArgs, 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 labels["dp_rank"] = dp_rank
self.metrics_collector = SchedulerMetricsCollector(labels=labels) 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]): def init_kv_events(self: Scheduler, kv_events_config: Optional[str]):
if self.enable_kv_cache_events: if self.enable_kv_cache_events:
self.kv_event_publisher = EventPublisherFactory.create( self.kv_event_publisher = EventPublisherFactory.create(
python/sglang/srt/server_args.py
+146 -171
View File
@@ -26,8 +26,6 @@ import random
import tempfile import tempfile
from typing import Any, Callable, Dict, List, Literal, Optional, Union from typing import Any, Callable, Dict, List, Literal, Optional, Union
import orjson
from sglang.srt.connector import ConnectorType from sglang.srt.connector import ConnectorType
from sglang.srt.environ import ToolStrictLevel, envs from sglang.srt.environ import ToolStrictLevel, envs
from sglang.srt.function_call.function_call_parser import FunctionCallParser from sglang.srt.function_call.function_call_parser import FunctionCallParser
@@ -65,7 +63,7 @@ from sglang.srt.utils.common import (
wait_port_available, wait_port_available,
xpu_has_xmx_support, 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 from sglang.utils import is_in_ci
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -190,7 +188,7 @@ FP8_GEMM_RUNNER_BACKEND_CHOICES = [
MAMBA_SSM_DTYPE_CHOICES = ["float32", "bfloat16"] 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 # Allow external code to add more choices
@@ -278,12 +276,6 @@ class ServerArgs:
nccl_port: Optional[int] = None nccl_port: Optional[int] = None
checkpoint_engine_wait_weights_before_ready: bool = False 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 # Quantization and data type
dtype: str = "auto" dtype: str = "auto"
quantization: Optional[str] = None quantization: Optional[str] = None
@@ -609,6 +601,12 @@ class ServerArgs:
# FIXME: hack to reduce ITL when decode bs is small # FIXME: hack to reduce ITL when decode bs is small
disaggregation_decode_polling_interval: int = 1 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 # For model weight update and weight loading
custom_weight_loader: Optional[List[str]] = None custom_weight_loader: Optional[List[str]] = None
weight_loader_disable_mmap: bool = False weight_loader_disable_mmap: bool = False
@@ -686,6 +684,7 @@ class ServerArgs:
self._handle_a2a_moe() self._handle_a2a_moe()
self._handle_eplb_and_dispatch() self._handle_eplb_and_dispatch()
self._handle_expert_distribution_metrics() self._handle_expert_distribution_metrics()
self._handle_elastic_ep()
# Handle pipeline parallelism. # Handle pipeline parallelism.
self._handle_pipeline_parallelism() self._handle_pipeline_parallelism()
@@ -693,9 +692,6 @@ class ServerArgs:
# Handle speculative decoding logic. # Handle speculative decoding logic.
self._handle_speculative_decoding() self._handle_speculative_decoding()
# Handle remote instance weight loader.
self._handle_remote_instance_weight_loader_start_seed_via_transfer_engine()
# Handle model loading format. # Handle model loading format.
self._handle_load_format() self._handle_load_format()
@@ -714,24 +710,15 @@ class ServerArgs:
# Validate cache settings. # Validate cache settings.
self._handle_cache_compatibility() self._handle_cache_compatibility()
# Validate metrics labels.
self._handle_metrics_labels()
# Handle deterministic inference. # Handle deterministic inference.
self._handle_deterministic_inference() self._handle_deterministic_inference()
# Handle exporting request-level metrics.
self._handle_request_metrics_exporters()
# Handle diffusion LLM inference. # Handle diffusion LLM inference.
self._handle_dllm_inference() self._handle_dllm_inference()
# Handle any other necessary validations. # Handle any other necessary validations.
self._handle_other_validations() self._handle_other_validations()
# Handle elastic expert parallelism.
self._handle_elastic_ep()
def _handle_deprecated_args(self): def _handle_deprecated_args(self):
# Handle deprecated tool call parsers # Handle deprecated tool call parsers
deprecated_tool_call_parsers = {"qwen25": "qwen", "glm45": "glm"} deprecated_tool_call_parsers = {"qwen25": "qwen", "glm45": "glm"}
@@ -752,17 +739,6 @@ class ServerArgs:
self.random_seed = random.randint(0, 1 << 30) self.random_seed = random.randint(0, 1 << 30)
if self.mm_process_config is None: if self.mm_process_config is None:
self.mm_process_config = {} 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 get_bool_env_var("SGLANG_USE_MODELSCOPE"):
@@ -774,6 +750,44 @@ class ServerArgs:
self.tokenizer_path, ignore_patterns=["*.bin", "*.safetensors"] 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): def _handle_gpu_memory_settings(self, gpu_mem):
""" """
Configure GPU memory-dependent settings including Configure GPU memory-dependent settings including
@@ -973,37 +987,13 @@ class ServerArgs:
return capture_sizes 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): def _handle_model_specific_adjustments(self):
from sglang.srt.configs.model_config import is_deepseek_nsa from sglang.srt.configs.model_config import is_deepseek_nsa
if parse_connector_type(self.model_path) == ConnectorType.INSTANCE: if parse_connector_type(self.model_path) == ConnectorType.INSTANCE:
return return
hf_config = self.get_hf_config() hf_config = self.get_model_config().hf_config
model_arch = hf_config.architectures[0] model_arch = hf_config.architectures[0]
if model_arch in [ if model_arch in [
@@ -1017,16 +1007,13 @@ class ServerArgs:
"MistralLarge3ForCausalLM", "MistralLarge3ForCausalLM",
"PixtralForConditionalGeneration", "PixtralForConditionalGeneration",
]: ]:
if is_deepseek_nsa(hf_config): # Set attention backend for DeepSeek
if ( if is_deepseek_nsa(hf_config): # DeepSeek 3.2
self.attention_backend is None if self.is_attention_backend_not_set():
and self.prefill_attention_backend is None
and self.decode_attention_backend is None
):
self.attention_backend = "nsa" 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 self.enable_dp_attention = True
logger.warning("DP attention is enabled for DeepSeek NSA.") logger.warning("DP attention is enabled for DeepSeek NSA.")
if self.enable_nsa_prefill_context_parallel: if self.enable_nsa_prefill_context_parallel:
@@ -1090,10 +1077,11 @@ class ServerArgs:
print_nsa_bool_env_vars() print_nsa_bool_env_vars()
else: else:
# DeepSeek V3/R1/V3.1
if self.enable_piecewise_cuda_graph: if self.enable_piecewise_cuda_graph:
logger.info("Piecewise CUDA graph is enabled, use MLA for prefill.") logger.info("Piecewise CUDA graph is enabled, use MLA for prefill.")
if is_cuda() and is_sm100_supported(): if is_sm100_supported():
if ( if (
self.attention_backend is None self.attention_backend is None
and self.prefill_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" "Use trtllm_mla as attention backend on sm100 for DeepseekV3ForCausalLM"
) )
# common to all Deepseek MoE models # Set moe backend for DeepSeek
if is_cuda() and is_sm100_supported(): if is_sm100_supported():
quantization_config = getattr(hf_config, "quantization_config", None) quantization_config = getattr(hf_config, "quantization_config", None)
quant_method = ( quant_method = (
quantization_config.get("quant_method") quantization_config.get("quant_method")
@@ -1116,7 +1104,7 @@ class ServerArgs:
# Default DeepSeek V3/R1 native FP8 when not explicitly set, # Default DeepSeek V3/R1 native FP8 when not explicitly set,
# Because we need this condition for an assertion in # Because we need this condition for an assertion in
# flashinfer_trtllm MoE runner backend. # 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" self.quantization = "fp8"
logger.info( logger.info(
"Quantization not specified, default to fp8 for DeepSeek on sm100" "Quantization not specified, default to fp8 for DeepSeek on sm100"
@@ -1134,14 +1122,11 @@ class ServerArgs:
) )
elif model_arch in ["GptOssForCausalLM"]: elif model_arch in ["GptOssForCausalLM"]:
if ( # Set attention backend for GPT-OSS
self.attention_backend is None if self.is_attention_backend_not_set():
and self.prefill_attention_backend is None if is_sm100_supported():
and self.decode_attention_backend is None
):
if is_cuda() and is_sm100_supported():
self.attention_backend = "trtllm_mha" self.attention_backend = "trtllm_mha"
elif is_cuda() and is_sm90_supported(): elif is_sm90_supported():
self.attention_backend = "fa3" self.attention_backend = "fa3"
else: else:
self.attention_backend = "triton" self.attention_backend = "triton"
@@ -1381,11 +1366,6 @@ class ServerArgs:
FLA_CHUNK_SIZE % self.page_size == 0 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}" ), 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: elif not self.disable_radix_cache:
logger.warning( logger.warning(
"Disabling overlap schedule since MambaRadixCache no_buffer is not compatible with " "Disabling overlap schedule since MambaRadixCache no_buffer is not compatible with "
@@ -1421,6 +1401,10 @@ class ServerArgs:
self.disable_radix_cache = True self.disable_radix_cache = True
self.disable_overlap_schedule = False 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 # TRTLLM AllReduce Fusion supports SM90/100/120, enable it by default
# for models with explicit support (DeepseekV3, GptOss, Glm4Moe, Qwen3Moe) # 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 # 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" and self.moe_a2a_backend == "none"
): ):
self.enable_flashinfer_allreduce_fusion = True self.enable_flashinfer_allreduce_fusion = True
logger.info(
f"Enable FlashInfer AllReduce Fusion by default for {model_arch}"
)
def _handle_sampling_backend(self): def _handle_sampling_backend(self):
if self.sampling_backend is None: if self.sampling_backend is None:
@@ -1520,7 +1501,7 @@ class ServerArgs:
self.attention_backend = "triton" self.attention_backend = "triton"
logger.warning( 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 # Torch native and flex attention backends
@@ -1977,7 +1958,7 @@ class ServerArgs:
"eagle speculative decoding." "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 [ if model_arch in [
"DeepseekV32ForCausalLM", "DeepseekV32ForCausalLM",
"DeepseekV3ForCausalLM", "DeepseekV3ForCausalLM",
@@ -2117,27 +2098,10 @@ class ServerArgs:
) )
self.load_format = "auto" self.load_format = "auto"
def _handle_encoder_disaggregation(self): # Check whether TransferEngine can be used when users want to start seed service that supports TransferEngine backend.
if self.enable_prefix_mm_cache and not self.encoder_only: if self.remote_instance_weight_loader_start_seed_via_transfer_engine:
raise ValueError( self.remote_instance_weight_loader_start_seed_via_transfer_engine = (
"--enable-prefix-mm-cache requires --encoder-only to be enabled" self.validate_transfer_engine()
)
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_pd_disaggregation(self): 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." "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): def _handle_tokenizer_batching(self):
if self.enable_tokenizer_batch_encode and self.enable_dynamic_batch_tokenizer: if self.enable_tokenizer_batch_encode and self.enable_dynamic_batch_tokenizer:
raise ValueError( raise ValueError(
@@ -2238,15 +2225,6 @@ class ServerArgs:
"Spec v2 and decode offload kv cache are incompatible and cannot be enabled together." "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): def _handle_deterministic_inference(self):
if self.rl_on_policy_target is not None: if self.rl_on_policy_target is not None:
logger.warning( logger.warning(
@@ -2265,7 +2243,7 @@ class ServerArgs:
is_deepseek_model = False is_deepseek_model = False
if parse_connector_type(self.model_path) != ConnectorType.INSTANCE: if parse_connector_type(self.model_path) != ConnectorType.INSTANCE:
try: try:
hf_config = self.get_hf_config() hf_config = self.get_model_config().hf_config
model_arch = hf_config.architectures[0] model_arch = hf_config.architectures[0]
is_deepseek_model = model_arch in [ is_deepseek_model = model_arch in [
"DeepseekV2ForCausalLM", "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." "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): def _handle_dllm_inference(self):
if self.dllm_algorithm is None: if self.dllm_algorithm is None:
return return
@@ -2372,13 +2343,6 @@ class ServerArgs:
self.disable_cuda_graph = True self.disable_cuda_graph = True
self.skip_server_warmup = 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 @staticmethod
def add_cli_args(parser: argparse.ArgumentParser): def add_cli_args(parser: argparse.ArgumentParser):
@@ -2538,32 +2502,6 @@ class ServerArgs:
"before serving inference requests.", "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 # Quantization and data type
parser.add_argument( parser.add_argument(
"--dtype", "--dtype",
@@ -3617,7 +3555,7 @@ class ServerArgs:
parser.add_argument( parser.add_argument(
"--mamba-scheduler-strategy", "--mamba-scheduler-strategy",
type=str, type=str,
choices=mamba_scheduler_strategy_CHOICES, choices=MAMBA_SCHEDULER_STRATEGY_CHOICES,
default=ServerArgs.mamba_scheduler_strategy, default=ServerArgs.mamba_scheduler_strategy,
help="The strategy to use for mamba radix cache.", 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.", 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 # Custom weight loader
parser.add_argument( parser.add_argument(
"--custom-weight-loader", "--custom-weight-loader",
@@ -4391,17 +4355,6 @@ class ServerArgs:
else: else:
return f"http://{self.host}:{self.port}" 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): def get_model_config(self):
# Lazy init to avoid circular import # Lazy init to avoid circular import
from sglang.srt.configs.model_config import ModelConfig from sglang.srt.configs.model_config import ModelConfig
@@ -4430,6 +4383,13 @@ class ServerArgs:
model_config = self.get_model_config() model_config = self.get_model_config()
return model_config.attention_arch == AttentionArch.MLA 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: def enable_mamba_extra_buffer(self) -> bool:
return self.mamba_scheduler_strategy == "extra_buffer" return self.mamba_scheduler_strategy == "extra_buffer"
@@ -4548,6 +4508,21 @@ class ServerArgs:
if self.model_impl == "mindspore": if self.model_impl == "mindspore":
assert is_npu(), "MindSpore model impl is only supported on Ascend npu." 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): def check_torch_2_9_1_cudnn_compatibility(self):
if get_bool_env_var("SGLANG_DISABLE_CUDNN_CHECK"): if get_bool_env_var("SGLANG_DISABLE_CUDNN_CHECK"):
return 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 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] arch = hf_config.architectures[0]
if self.speculative_algorithm == "STANDALONE": if self.speculative_algorithm == "STANDALONE":
# The default value for standalone speculative decoding # The default value for standalone speculative decoding
python/sglang/srt/utils/common.py
+98 -107
View File
@@ -49,7 +49,7 @@ import warnings
from collections import OrderedDict, defaultdict from collections import OrderedDict, defaultdict
from contextlib import contextmanager from contextlib import contextmanager
from dataclasses import dataclass from dataclasses import dataclass
from functools import lru_cache from functools import lru_cache, partial
from importlib.metadata import PackageNotFoundError, version from importlib.metadata import PackageNotFoundError, version
from importlib.util import find_spec from importlib.util import find_spec
from io import BytesIO from io import BytesIO
@@ -105,22 +105,6 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__) 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 # https://pytorch.org/docs/stable/notes/hip.html#checking-for-hip
@lru_cache(maxsize=1) @lru_cache(maxsize=1)
@@ -129,6 +113,7 @@ def is_hip() -> bool:
if is_hip(): if is_hip():
HIP_FP8_E4M3_FNUZ_MAX = 224.0
FP8_E4M3_MAX = HIP_FP8_E4M3_FNUZ_MAX FP8_E4M3_MAX = HIP_FP8_E4M3_FNUZ_MAX
else: else:
FP8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max FP8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
@@ -191,14 +176,6 @@ def get_cuda_version():
return (0, 0) 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 @contextmanager
def device_context(device: torch.device): def device_context(device: torch.device):
if device.type == "cpu" and is_cpu(): if device.type == "cpu" and is_cpu():
@@ -213,96 +190,49 @@ def device_context(device: torch.device):
raise ValueError(f"Unknown device module: {device}") raise ValueError(f"Unknown device module: {device}")
is_ampere_with_cuda_12_3 = lambda: _check(8) def _check_cuda_device_version(
is_hopper_with_cuda_12_3 = lambda: _check(9) device_capability_majors: List[int], cuda_version: Tuple[int, int]
):
@lru_cache(maxsize=1)
def is_blackwell():
if not is_cuda(): if not is_cuda():
return False return False
return torch.cuda.get_device_capability()[0] in [10, 12] 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_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"
) )
@lru_cache(maxsize=1) is_ampere_with_cuda_12_3 = lru_cache(maxsize=1)(
def is_sm100_supported(device=None) -> bool: partial(
if not is_cuda(): _check_cuda_device_version, device_capability_majors=[8], cuda_version=(12, 3)
return False
return (torch.cuda.get_device_capability(device)[0] == 10) and (
torch.version.cuda >= "12.8"
) )
)
is_hopper_with_cuda_12_3 = lru_cache(maxsize=1)(
@lru_cache(maxsize=1) partial(
def is_sm90_supported(device=None) -> bool: _check_cuda_device_version, device_capability_majors=[9], cuda_version=(12, 3)
if not is_cuda():
return False
return (torch.cuda.get_device_capability(device)[0] == 9) and (
torch.version.cuda >= "12.3"
) )
)
is_blackwell_supported = is_blackwell = lru_cache(maxsize=1)(
_warned_bool_env_var_keys = set() partial(
_check_cuda_device_version,
device_capability_majors=[10, 12],
def get_bool_env_var(name: str, default: str = "false") -> bool: cuda_version=(12, 8),
# 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) )
is_sm120_supported = lru_cache(maxsize=1)(
return value in truthy_values partial(
_check_cuda_device_version, device_capability_majors=[12], cuda_version=(12, 8)
)
def get_int_env_var(name: str, default: int = 0) -> int: )
# FIXME: move your environment variable to sglang.srt.environ is_sm100_supported = lru_cache(maxsize=1)(
value = os.getenv(name) partial(
if value is None or not value.strip(): _check_cuda_device_version, device_capability_majors=[10], cuda_version=(12, 8)
return default )
try: )
return int(value) is_sm90_supported = lru_cache(maxsize=1)(
except ValueError: partial(
return default _check_cuda_device_version, device_capability_majors=[9], cuda_version=(12, 3)
)
)
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"]
try: try:
@@ -364,6 +294,53 @@ def random_uuid() -> str:
return str(uuid.uuid4().hex) 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( _ENABLE_TORCH_INFERENCE_MODE = get_bool_env_var(
"SGLANG_ENABLE_TORCH_INFERENCE_MODE", "false" "SGLANG_ENABLE_TORCH_INFERENCE_MODE", "false"
) )
@@ -421,6 +398,10 @@ class DynamicGradMode(_DecoratorContextManager):
return self.__class__() return self.__class__()
show_time_cost = False
time_infos = {}
def enable_show_time_cost(): def enable_show_time_cost():
global show_time_cost global show_time_cost
show_time_cost = True 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: if count % log_interval == 0:
logger.warning(message) logger.warning(message)
setattr(obj, counter_name, count + 1) 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
test/srt/openai_server/basic/test_serving_chat.py
+2 -4
View File
@@ -36,7 +36,7 @@ class _MockTokenizerManager:
# Mock hf_config for _use_dpsk_v32_encoding check # Mock hf_config for _use_dpsk_v32_encoding check
mock_hf_config = Mock() mock_hf_config = Mock()
mock_hf_config.architectures = ["LlamaForCausalLM"] 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" self.chat_template_name: Optional[str] = "llama-3"
@@ -603,9 +603,7 @@ class ServingChatTestCase(unittest.TestCase):
mock_hf_config = Mock() mock_hf_config = Mock()
mock_hf_config.architectures = ["DeepseekV32ForCausalLM"] mock_hf_config.architectures = ["DeepseekV32ForCausalLM"]
tokenizer_manager.server_args.get_hf_config = Mock( tokenizer_manager.model_config.hf_config = mock_hf_config
return_value=mock_hf_config
)
# Case 1: No chat template in tokenizer -> should use dpsk encoding # Case 1: No chat template in tokenizer -> should use dpsk encoding
tokenizer_manager.tokenizer = Mock() tokenizer_manager.tokenizer = Mock()