From 7b36c47b3be2037556df44d83ad279263accd56e Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Sat, 25 Oct 2025 23:50:12 -0700 Subject: [PATCH] Clean up attention backend selection code & Other minor rename (#12136) --- python/sglang/srt/entrypoints/http_server.py | 5 + python/sglang/srt/managers/scheduler.py | 21 +- .../srt/managers/scheduler_metrics_mixin.py | 27 ++- .../sglang/srt/model_executor/model_runner.py | 145 ------------ python/sglang/srt/server_args.py | 208 +++++++++++++++--- python/sglang/srt/utils/common.py | 151 ++++++------- 6 files changed, 282 insertions(+), 275 deletions(-) diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 057317079..ee2586fde 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -498,6 +498,11 @@ async def get_server_info(): internal_states: List[Dict[Any, Any]] = ( await _global_state.tokenizer_manager.get_internal_state() ) + + # This field is not serializable. + if hasattr(_global_state.tokenizer_manager.server_args, "model_config"): + del _global_state.tokenizer_manager.server_args.model_config + return { **dataclasses.asdict(_global_state.tokenizer_manager.server_args), **_global_state.scheduler_info, diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index c1ae06612..4d04eda83 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2325,10 +2325,10 @@ class Scheduler( self.num_generated_tokens = 0 self.forward_ct_decode = 0 - self.spec_num_total_accepted_tokens = 0 - self.spec_num_total_forward_ct = 0 - self.cum_spec_accept_length = 0 - self.cum_spec_accept_count = 0 + self.spec_num_accepted_tokens = 0 + self.spec_num_forward_ct = 0 + self.spec_total_num_accepted_tokens = 0 + self.spec_total_num_forward_ct = 0 torch.cuda.empty_cache() logger.info("Cache flushed successfully!") if_success = True @@ -2401,13 +2401,16 @@ class Scheduler( self.tp_worker.model_runner.graph_mem_usage, 2 ) - if not self.spec_algorithm.is_none() and self.cum_spec_accept_count > 0: + if not self.spec_algorithm.is_none() and self.spec_total_num_forward_ct > 0: ret["avg_spec_accept_length"] = ( - self.cum_spec_accept_length / self.cum_spec_accept_count + self.spec_total_num_accepted_tokens / self.spec_total_num_forward_ct ) if RECORD_STEP_TIME: ret["step_time_dict"] = self.step_time_dict + # This field is not serializable. + ret.pop("model_config", None) + return GetInternalStateReqOutput(internal_state=ret) def set_internal_state(self, recv_req: SetInternalStateReq): @@ -2434,12 +2437,12 @@ class Scheduler( if_success = False break if if_success: - if not self.spec_algorithm.is_none() and self.cum_spec_accept_count > 0: + if not self.spec_algorithm.is_none() and self.spec_total_num_forward_ct > 0: avg_spec_accept_length = ( - self.cum_spec_accept_length / self.cum_spec_accept_count + self.spec_total_num_accepted_tokens / self.spec_total_num_forward_ct ) logger.info(f"{avg_spec_accept_length=}") - self.cum_spec_accept_length = self.cum_spec_accept_count = 0 + self.spec_total_num_accepted_tokens = self.spec_total_num_forward_ct = 0 for k, v in server_args_dict.items(): setattr(get_global_server_args(), k, v) logger.info(f"Global server args updated! {get_global_server_args()=}") diff --git a/python/sglang/srt/managers/scheduler_metrics_mixin.py b/python/sglang/srt/managers/scheduler_metrics_mixin.py index aa2ae130d..783ff4627 100644 --- a/python/sglang/srt/managers/scheduler_metrics_mixin.py +++ b/python/sglang/srt/managers/scheduler_metrics_mixin.py @@ -39,10 +39,13 @@ class SchedulerMetricsMixin: self.last_gen_throughput: float = 0.0 self.last_input_throughput: float = 0.0 self.step_time_dict = defaultdict(list) # Dict[batch size -> step time] - self.spec_num_total_accepted_tokens = 0 - self.spec_num_total_forward_ct = 0 - self.cum_spec_accept_length = 0 - self.cum_spec_accept_count = 0 + + # The number of accepted tokens and forward ct for the recent `decode_log_interval` batches (for logging) + self.spec_num_accepted_tokens = 0 + self.spec_num_forward_ct = 0 + # The total number of accepted tokens and forward ct for the whole server lifetime + self.spec_total_num_accepted_tokens = 0 + self.spec_total_num_forward_ct = 0 self.kv_transfer_speed_gb_s: float = 0.0 self.kv_transfer_latency_ms: float = 0.0 @@ -67,8 +70,8 @@ class SchedulerMetricsMixin: ) def update_spec_metrics(self: Scheduler, bs: int, num_accepted_tokens: int): - self.spec_num_total_accepted_tokens += num_accepted_tokens + bs - self.spec_num_total_forward_ct += bs + self.spec_num_accepted_tokens += num_accepted_tokens + bs + self.spec_num_forward_ct += bs self.num_generated_tokens += num_accepted_tokens def log_prefill_stats( @@ -253,20 +256,20 @@ class SchedulerMetricsMixin: spec_accept_rate = 0 else: spec_accept_length = ( - self.spec_num_total_accepted_tokens / self.spec_num_total_forward_ct + self.spec_num_accepted_tokens / self.spec_num_forward_ct ) # Calculate acceptance rate: accepted tokens / total draft tokens - total_draft_tokens = self.spec_num_total_forward_ct * ( + total_draft_tokens = self.spec_num_forward_ct * ( (self.server_args.speculative_num_steps or 0) + 1 ) spec_accept_rate = ( - self.spec_num_total_accepted_tokens / total_draft_tokens + self.spec_num_accepted_tokens / total_draft_tokens if total_draft_tokens > 0 else 0 ) - self.cum_spec_accept_length += self.spec_num_total_accepted_tokens - self.cum_spec_accept_count += self.spec_num_total_forward_ct - self.spec_num_total_accepted_tokens = self.spec_num_total_forward_ct = 0 + self.spec_total_num_accepted_tokens += self.spec_num_accepted_tokens + self.spec_total_num_forward_ct += self.spec_num_forward_ct + self.spec_num_accepted_tokens = self.spec_num_forward_ct = 0 msg += f"accept len: {spec_accept_length:.2f}, accept rate: {spec_accept_rate:.2f}, " cache_hit_rate = 0.0 diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index ade621ddc..eac8f4d5e 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -131,13 +131,8 @@ from sglang.srt.utils import ( get_bool_env_var, get_cpu_ids_by_node, init_custom_process_group, - is_fa3_default_architecture, - is_flashinfer_available, is_hip, - is_hopper_with_cuda_12_3, - is_no_spec_infer_or_topk_one, is_npu, - is_sm100_supported, log_info_on_rank0, monkey_patch_p2p_access_check, set_cuda_arch, @@ -502,121 +497,6 @@ class ModelRunner: def model_specific_adjustment(self): server_args = self.server_args - if ( - server_args.attention_backend == "intel_amx" - and server_args.device == "cpu" - and not _is_cpu_amx_available - ): - logger.info( - "The current platform does not support Intel AMX, will fallback to torch_native backend." - ) - server_args.attention_backend = "torch_native" - - if ( - server_args.attention_backend == "intel_xpu" - and server_args.device == "xpu" - and not _is_xpu_xmx_available - ): - logger.info( - "The current platform does not support Intel XMX, will fallback to triton backend." - ) - server_args.attention_backend = "triton" - - if server_args.prefill_attention_backend is not None and ( - server_args.prefill_attention_backend - == server_args.decode_attention_backend - ): # override the default attention backend - server_args.attention_backend = server_args.prefill_attention_backend - - if ( - getattr(self.model_config.hf_config, "dual_chunk_attention_config", None) - is not None - ): - if server_args.attention_backend is None: - server_args.attention_backend = "dual_chunk_flash_attn" - logger.info("Dual chunk attention is turned on by default.") - elif server_args.attention_backend != "dual_chunk_flash_attn": - raise ValueError( - "Dual chunk attention is enabled, but attention backend is set to " - f"{server_args.attention_backend}. Please set it to 'dual_chunk_flash_attn'." - ) - - if server_args.attention_backend is None: - """ - Auto select the fastest attention backend. - - 1. Models with MHA Architecture (e.g: Llama, QWen) - 1.1 We will turn on FA3 on hopper unless user use spec decode with topk > 1 or page_size > 1. - 1.2 In other cases, we will use flashinfer if available, otherwise use triton. - 2. Models with MLA Architecture and using FA3 - 2.1 We will use FA3 backend on hopper. - 2.2 We will use Flashinfer backend on blackwell. - 2.3 Otherwise, we will use triton backend. - """ - - if not self.use_mla_backend: - # MHA architecture - if ( - is_hopper_with_cuda_12_3() - and is_no_spec_infer_or_topk_one(server_args) - and is_fa3_default_architecture(self.model_config.hf_config) - ): - server_args.attention_backend = "fa3" - elif _is_hip: - server_args.attention_backend = "aiter" - elif _is_npu: - server_args.attention_backend = "ascend" - else: - server_args.attention_backend = ( - "flashinfer" if is_flashinfer_available() else "triton" - ) - else: - # MLA architecture - if is_hopper_with_cuda_12_3(): - server_args.attention_backend = "fa3" - elif is_sm100_supported(): - server_args.attention_backend = "flashinfer" - elif _is_hip: - head_num = self.model_config.get_num_kv_heads(self.tp_size) - # TODO current aiter only support head number 16 or 128 head number - if head_num == 128 or head_num == 16: - server_args.attention_backend = "aiter" - else: - server_args.attention_backend = "triton" - elif _is_npu: - server_args.attention_backend = "ascend" - else: - server_args.attention_backend = "triton" - log_info_on_rank0( - logger, - f"Attention backend not explicitly specified. Use {server_args.attention_backend} backend by default.", - ) - elif self.use_mla_backend: - if server_args.device != "cpu": - if server_args.attention_backend in MLA_ATTENTION_BACKENDS: - logger.info( - f"MLA optimization is turned on. Use {server_args.attention_backend} backend." - ) - else: - raise ValueError( - f"Invalid attention backend for MLA: {server_args.attention_backend}" - ) - else: - if server_args.attention_backend != "intel_amx": - raise ValueError( - "MLA optimization not supported on CPU except for intel_amx backend." - ) - - if ( - server_args.attention_backend == "fa3" - and server_args.kv_cache_dtype == "fp8_e5m2" - ): - logger.warning( - "FlashAttention3 only supports fp8_e4m3 if using FP8; " - "Setting attention backend to triton." - ) - server_args.attention_backend = "triton" - if server_args.enable_double_sparsity: logger.info( "Double sparsity optimization is turned on. Use triton backend without CUDA graph." @@ -642,31 +522,6 @@ class ModelRunner: if not server_args.disable_chunked_prefix_cache: log_info_on_rank0(logger, "Chunked prefix cache is turned on.") - if server_args.attention_backend == "aiter": - if self.model_config.context_len > 8192: - self.mem_fraction_static *= 0.85 - - if ( - server_args.enable_hierarchical_cache - and server_args.hicache_io_backend == "kernel" - ): - # fix for the compatibility issue with FlashAttention3 decoding and HiCache kernel backend - if server_args.decode_attention_backend is None: - if not self.use_mla_backend: - server_args.decode_attention_backend = ( - "flashinfer" if is_flashinfer_available() else "triton" - ) - else: - server_args.decode_attention_backend = ( - "flashinfer" if is_sm100_supported() else "triton" - ) - elif server_args.decode_attention_backend == "fa3": - server_args.hicache_io_backend = "direct" - logger.warning( - "FlashAttention3 decode backend is not compatible with hierarchical cache. " - "Setting hicache_io_backend to vanilla I/O, which may lead to suboptimal performance with small page sizes." - ) - if self.model_config.hf_config.model_type == "qwen3_vl_moe": if ( quantization_config := getattr( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index d2078e473..db8a9169b 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -34,12 +34,16 @@ from sglang.srt.utils.common import ( LORA_TARGET_ALL_MODULES, SUPPORTED_LORA_TARGET_MODULES, configure_ipv6, + cpu_has_amx_support, get_device, get_device_memory_capacity, get_device_sm, is_cuda, + is_fa3_default_architecture, is_flashinfer_available, is_hip, + is_hopper_with_cuda_12_3, + is_no_spec_infer_or_topk_one, is_npu, is_port_available, is_remote_url, @@ -51,6 +55,7 @@ from sglang.srt.utils.common import ( json_list_type, nullable_str, parse_connector_type, + xpu_has_xmx_support, ) from sglang.srt.utils.hf_transformers_utils import check_gguf_file, get_config from sglang.utils import is_in_ci @@ -545,6 +550,9 @@ class ServerArgs: # Apply model-specific adjustments. self._handle_model_specific_adjustments() + # Handle Hicache settings. + self._handle_hicache() + # Set kernel backends. self._handle_sampling_backend() self._handle_attention_backend_compatibility() @@ -567,9 +575,6 @@ class ServerArgs: # Handle pipeline parallelism. self._handle_pipeline_parallelism() - # Handle Hicache settings. - self._handle_hicache() - # Handle speculative decoding logic. self._handle_speculative_decoding() @@ -779,11 +784,9 @@ class ServerArgs: else 0.88 ) - # Lazy init to avoid circular import - # Multimodal models need more memory for the image processor - from sglang.srt.configs.model_config import ModelConfig - - model_config = ModelConfig.from_server_args(self) + # Multimodal models need more memory for the image processing, + # so we adjust the mem_fraction_static accordingly. + model_config = self.get_model_config() if model_config.is_multimodal: self.adjust_mem_fraction_for_vlm(model_config) @@ -1042,6 +1045,67 @@ class ServerArgs: ) def _handle_attention_backend_compatibility(self): + model_config = self.get_model_config() + use_mla_backend = self.use_mla_backend() + + if self.prefill_attention_backend is not None and ( + self.prefill_attention_backend == self.decode_attention_backend + ): # override the default attention backend + self.attention_backend = self.prefill_attention_backend + + # Pick the default attention backend if not specified + if self.attention_backend is None: + """ + Auto select the fastest attention backend. + + 1. Models with MHA Architecture (e.g: Llama, QWen) + 1.1 We will turn on FA3 on hopper unless user use spec decode with topk > 1 or page_size > 1. + 1.2 In other cases, we will use flashinfer if available, otherwise use triton. + 2. Models with MLA Architecture and using FA3 + 2.1 We will use FA3 backend on hopper. + 2.2 We will use Flashinfer backend on blackwell. + 2.3 Otherwise, we will use triton backend. + """ + + if not use_mla_backend: + # MHA architecture + if ( + is_hopper_with_cuda_12_3() + and is_no_spec_infer_or_topk_one(self) + and is_fa3_default_architecture(self.model_config.hf_config) + ): + self.attention_backend = "fa3" + elif is_hip(): + self.attention_backend = "aiter" + elif is_npu(): + self.attention_backend = "ascend" + else: + self.attention_backend = ( + "flashinfer" if is_flashinfer_available() else "triton" + ) + else: + # MLA architecture + if is_hopper_with_cuda_12_3(): + self.attention_backend = "fa3" + elif is_sm100_supported(): + self.attention_backend = "flashinfer" + elif is_hip(): + head_num = model_config.get_num_kv_heads(self.tp_size) + # TODO current aiter only support head number 16 or 128 head number + if head_num == 128 or head_num == 16: + self.attention_backend = "aiter" + else: + self.attention_backend = "triton" + elif is_npu(): + self.attention_backend = "ascend" + else: + self.attention_backend = "triton" + + logger.warning( + f"Attention backend not explicitly specified. Use {self.attention_backend} backend by default." + ) + + # Torch native and flex attention backends if self.attention_backend == "torch_native": logger.warning( "Cuda graph is disabled because of using torch native attention backend" @@ -1057,12 +1121,7 @@ class ServerArgs: self.speculative_algorithm is None ), "Speculative decoding is currently not supported with Flex Attention backend" - if is_npu() and self.attention_backend in ["ascend"]: - logger.warning( - "At this moment Ascend attention backend only supports a page_size of 128, change page_size to 128." - ) - self.page_size = 128 - + # Major NVIDIA platforms backends if ( self.attention_backend == "flashmla" or self.decode_attention_backend == "flashmla" @@ -1117,19 +1176,13 @@ class ServerArgs: ) self.page_size = 64 - if self.attention_backend == "dual_chunk_flash_attn": + if self.attention_backend == "fa3" and self.kv_cache_dtype == "fp8_e5m2": logger.warning( - "Mixed chunk and radix cache are disabled when using dual-chunk flash attention backend" + "FlashAttention3 only supports fp8_e4m3 if using FP8; " + "Setting attention backend to triton." ) - self.enable_mixed_chunk = False - self.disable_radix_cache = True + self.attention_backend = "triton" - if self.attention_backend == "intel_xpu": - if self.page_size not in [32, 64, 128]: - logger.warning( - f"Intel XPU attention backend only supports page_size of 32, 64 or 128, changing page_size from {self.page_size} to 128." - ) - self.page_size = 128 if self.attention_backend == "fa4" or self.decode_attention_backend == "fa4": raise ValueError( "FA4 backend is only supported for prefill. Please use `--prefill-attention-backend fa4` instead." @@ -1140,6 +1193,66 @@ class ServerArgs: ) self.page_size = 128 + # AMD platforms backends + if self.attention_backend == "aiter": + if model_config.context_len > 8192: + self.mem_fraction_static *= 0.90 + + # NPU platforms backends + if is_npu() and self.attention_backend in ["ascend"]: + logger.warning( + "At this moment Ascend attention backend only supports a page_size of 128, change page_size to 128." + ) + self.page_size = 128 + + # Other platforms backends + if ( + self.attention_backend == "intel_amx" + and self.device == "cpu" + and not cpu_has_amx_support() + ): + logger.warning( + "The current platform does not support Intel AMX, will fallback to torch_native backend." + ) + self.attention_backend = "torch_native" + + if ( + self.attention_backend == "intel_xpu" + and self.device == "xpu" + and not xpu_has_xmx_support() + ): + logger.warning( + "The current platform does not support Intel XMX, will fallback to triton backend." + ) + self.attention_backend = "triton" + + if self.attention_backend == "intel_xpu": + if self.page_size not in [32, 64, 128]: + logger.warning( + f"Intel XPU attention backend only supports page_size of 32, 64 or 128, changing page_size from {self.page_size} to 128." + ) + self.page_size = 128 + + # Dual chunk flash attention backend + if ( + getattr(model_config.hf_config, "dual_chunk_attention_config", None) + is not None + ): + if self.attention_backend is None: + self.attention_backend = "dual_chunk_flash_attn" + logger.info("Dual chunk attention is turned on by default.") + elif self.attention_backend != "dual_chunk_flash_attn": + raise ValueError( + "Dual chunk attention is enabled, but attention backend is set to " + f"{self.attention_backend}. Please set it to 'dual_chunk_flash_attn'." + ) + if self.attention_backend == "dual_chunk_flash_attn": + logger.warning( + "Mixed chunk and radix cache are disabled when using dual-chunk flash attention backend" + ) + self.enable_mixed_chunk = False + self.disable_radix_cache = True + def _handle_page_size(self): if self.page_size is None: self.page_size = 1 @@ -1283,6 +1396,24 @@ class ServerArgs: "Page first direct layout only support direct io backend" ) + if self.enable_hierarchical_cache and self.hicache_io_backend == "kernel": + # fix for the compatibility issue with FlashAttention3 decoding and HiCache kernel backend + if self.decode_attention_backend is None: + if not self.use_mla_backend(): + self.decode_attention_backend = ( + "flashinfer" if is_flashinfer_available() else "triton" + ) + else: + self.decode_attention_backend = ( + "flashinfer" if is_sm100_supported() else "triton" + ) + elif self.decode_attention_backend == "fa3": + self.hicache_io_backend = "direct" + logger.warning( + "FlashAttention3 decode backend is not compatible with hierarchical cache. " + "Setting hicache_io_backend to vanilla I/O, which may lead to suboptimal performance with small page sizes." + ) + def _handle_speculative_decoding(self): if self.speculative_algorithm == "NEXTN": self.speculative_algorithm = "EAGLE" @@ -3355,19 +3486,34 @@ class ServerArgs: ) return hf_config - def get_attention_backends(server_args): + def get_model_config(self): + # Lazy init to avoid circular import + from sglang.srt.configs.model_config import ModelConfig + + if hasattr(self, "model_config"): + return self.model_config + self.model_config = ModelConfig.from_server_args(self) + return self.model_config + + def get_attention_backends(self): prefill_attention_backend_str = ( - server_args.prefill_attention_backend - if server_args.prefill_attention_backend - else server_args.attention_backend + self.prefill_attention_backend + if self.prefill_attention_backend + else self.attention_backend ) decode_attention_backend_str = ( - server_args.decode_attention_backend - if server_args.decode_attention_backend - else server_args.attention_backend + self.decode_attention_backend + if self.decode_attention_backend + else self.attention_backend ) return prefill_attention_backend_str, decode_attention_backend_str + def use_mla_backend(self): + from sglang.srt.configs.model_config import AttentionArch + + model_config = self.get_model_config() + return model_config.attention_arch == AttentionArch.MLA + def check_server_args(self): # Check parallel size constraints assert ( diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index eb7e63e6a..8ea262bef 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -2096,80 +2096,80 @@ class MultiprocessingSerializer: # Decode base64 string to bytes data = pybase64.b64decode(data, validate=True) - class SafeUnpickler(pickle.Unpickler): - ALLOWED_MODULE_PREFIXES = { - # --- Python types --- - "builtins.", - "collections.", - "copyreg.", - "functools.", - "itertools.", - "operator.", - "types.", - "weakref.", - # --- PyTorch types --- - "torch.", - "torch._tensor.", - "torch.storage.", - "torch.nn.parameter.", - "torch.autograd.function.", - # --- torch distributed --- - "torch.distributed.", - "torch.distributed._shard.", - "torch.distributed._composable.", - "torch._C._distributed_c10d.", - "torch._C._distributed_fsdp.", - "torch.distributed.optim.", - # --- multiprocessing --- - "multiprocessing.resource_sharer.", - "multiprocessing.reduction.", - "pickletools.", - # --- PEFT / LoRA --- - "peft.", - "transformers.", - "huggingface_hub.", - # --- SGLang & Unitest --- - "sglang.srt.weight_sync.tensor_bucket.", - "sglang.srt.model_executor.model_runner.", - "sglang.srt.layers.", - "sglang.srt.utils.", - } - - DENY_CLASSES = { - ("builtins", "eval"), - ("builtins", "exec"), - ("builtins", "compile"), - ("os", "system"), - ("subprocess", "Popen"), - ("subprocess", "run"), - ("codecs", "decode"), - ("types", "CodeType"), - ("types", "FunctionType"), - } - - def find_class(self, module, name): - # Block deterministic attacks - if (module, name) in self.DENY_CLASSES: - raise RuntimeError( - f"Blocked unsafe class loading ({module}.{name}), " - f"to prevent exploitation of CVE-2025-10164" - ) - # Allowlist of safe-to-load modules. - if any( - (module + ".").startswith(prefix) - for prefix in self.ALLOWED_MODULE_PREFIXES - ): - return super().find_class(module, name) - - # Block everything else. (Potential attack surface) - raise RuntimeError( - f"Blocked unsafe class loading ({module}.{name}), " - f"to prevent exploitation of CVE-2025-10164" - ) - return SafeUnpickler(io.BytesIO(data)).load() +class SafeUnpickler(pickle.Unpickler): + ALLOWED_MODULE_PREFIXES = { + # --- Python types --- + "builtins.", + "collections.", + "copyreg.", + "functools.", + "itertools.", + "operator.", + "types.", + "weakref.", + # --- PyTorch types --- + "torch.", + "torch._tensor.", + "torch.storage.", + "torch.nn.parameter.", + "torch.autograd.function.", + # --- torch distributed --- + "torch.distributed.", + "torch.distributed._shard.", + "torch.distributed._composable.", + "torch._C._distributed_c10d.", + "torch._C._distributed_fsdp.", + "torch.distributed.optim.", + # --- multiprocessing --- + "multiprocessing.resource_sharer.", + "multiprocessing.reduction.", + "pickletools.", + # --- PEFT / LoRA --- + "peft.", + "transformers.", + "huggingface_hub.", + # --- SGLang & Unitest --- + "sglang.srt.weight_sync.tensor_bucket.", + "sglang.srt.model_executor.model_runner.", + "sglang.srt.layers.", + "sglang.srt.utils.", + } + + DENY_CLASSES = { + ("builtins", "eval"), + ("builtins", "exec"), + ("builtins", "compile"), + ("os", "system"), + ("subprocess", "Popen"), + ("subprocess", "run"), + ("codecs", "decode"), + ("types", "CodeType"), + ("types", "FunctionType"), + } + + def find_class(self, module, name): + # Block deterministic attacks + if (module, name) in self.DENY_CLASSES: + raise RuntimeError( + f"Blocked unsafe class loading ({module}.{name}), " + f"to prevent exploitation of CVE-2025-10164" + ) + # Allowlist of safe-to-load modules. + if any( + (module + ".").startswith(prefix) for prefix in self.ALLOWED_MODULE_PREFIXES + ): + return super().find_class(module, name) + + # Block everything else. (Potential attack surface) + raise RuntimeError( + f"Blocked unsafe class loading ({module}.{name}), " + f"to prevent exploitation of CVE-2025-10164" + ) + + def debug_timing(func): # todo: replace with a more organized instrumentation def wrapper(*args, **kwargs): @@ -2620,17 +2620,12 @@ def get_local_ip_auto(fallback: str = None) -> str: raise ValueError("Can not get local ip") -def is_page_size_one(server_args): - return server_args.page_size == 1 - - # TODO(hebiao064): Accelerate FA3 Spec Decode with topk > 1. # TODO(hebiao064): Improve the acc rate for FA3 Spec Decode with topk == 1 and page_size > 1. def is_no_spec_infer_or_topk_one(server_args): return server_args.speculative_eagle_topk is None or ( - server_args.speculative_eagle_topk is not None - and server_args.speculative_eagle_topk == 1 - and is_page_size_one(server_args) + server_args.speculative_eagle_topk == 1 + and (server_args.page_size == 1 or server_args.page_size is None) )