From f4e835af2faa18220e327264efaa02e74dc2f60a Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Thu, 25 Dec 2025 18:13:30 +0800 Subject: [PATCH] Cleanup `ModelRunner` (#15802) --- docs/references/environment_variables.md | 3 +- python/sglang/srt/environ.py | 7 +- .../sglang/srt/model_executor/model_runner.py | 116 ++++++++---------- .../spec/eagle/test_eagle_infer_b.py | 10 +- 4 files changed, 62 insertions(+), 74 deletions(-) diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index b53f3e1b2..c9861c445 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -67,7 +67,7 @@ SGLang supports various environment variables that can be used to configure its | `SGLANG_DEBUG_MEMORY_POOL` | Enable memory pool debugging | `false` | | `SGLANG_CLIP_MAX_NEW_TOKENS_ESTIMATION` | Clip max new tokens estimation for memory planning | `4096` | | `SGLANG_DETOKENIZER_MAX_STATES` | Maximum states for detokenizer | Default value based on system | -| `SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK` | Disable checks for memory imbalance across Tensor Parallel ranks | Not set (defaults to enabled check) | +| `SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK` | Enable checks for memory imbalance across Tensor Parallel ranks | `true` | ## Model-Specific Options @@ -112,7 +112,6 @@ SGLang supports various environment variables that can be used to configure its | `SGLANG_TEST_RETRACT_NO_PREFILL_BS` | When SGLANG_TEST_RETRACT is enabled, no prefill is performed if the batch size exceeds SGLANG_TEST_RETRACT_NO_PREFILL_BS. | `2 ** 31` | | `SGLANG_RECORD_STEP_TIME` | Record step time for profiling | `false` | | `SGLANG_TEST_REQUEST_TIME_STATS` | Test request time statistics | `false` | -| `SGLANG_CI_SMALL_KV_SIZE` | Use small KV cache size in CI | `-1` | ## Profiling & Benchmarking diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index fd001185c..c2cc903a4 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -146,6 +146,7 @@ class Envs: SGLANG_TEST_MAX_RETRY = EnvInt(None) # Test & Debug + SGLANG_DETECT_SLOW_RANK = EnvBool(False) SGLANG_TEST_STUCK_DETOKENIZER = EnvFloat(0) SGLANG_TEST_STUCK_DP_CONTROLLER = EnvFloat(0) SGLANG_TEST_STUCK_TOKENIZER = EnvFloat(0) @@ -166,6 +167,7 @@ class Envs: SGLANG_OTLP_EXPORTER_SCHEDULE_DELAY_MILLIS = EnvInt(500) SGLANG_OTLP_EXPORTER_MAX_EXPORT_BATCH_SIZE = EnvInt(64) SGLANG_NATIVE_MOVE_KV_CACHE = EnvBool(False) + SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(True) # Scheduler: memory leak test SGLANG_TEST_RETRACT = EnvBool(False) @@ -173,7 +175,6 @@ class Envs: SGLANG_TEST_RETRACT_NO_PREFILL_BS = EnvInt(2 ** 31) SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY = EnvInt(0) SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE = EnvBool(True) - SGLANG_CI_SMALL_KV_SIZE = EnvInt(-1) # Scheduler: new token ratio hyperparameters SGLANG_INIT_NEW_TOKEN_RATIO = EnvFloat(0.7) @@ -417,6 +418,10 @@ def _convert_SGL_to_SGLANG(): _print_deprecated_env( "SGLANG_MOE_NVFP4_DISPATCH", "SGLANG_CUTEDSL_MOE_NVFP4_DISPATCH" ) + _print_deprecated_env( + "SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK", + "SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK", + ) for key, value in os.environ.items(): if key.startswith("SGL_"): diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index f788dee80..5218e1649 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -164,11 +164,9 @@ from sglang.srt.utils import ( dynamic_import, enable_show_time_cost, get_available_gpu_memory, - get_bool_env_var, get_cpu_ids_by_node, get_local_ip_auto, init_custom_process_group, - is_cuda, is_float4_e2m1fn_x2, is_hip, is_npu, @@ -180,7 +178,6 @@ from sglang.srt.utils import ( reserve_rope_cache_for_long_sequences, set_cuda_arch, slow_rank_detector, - xpu_has_xmx_support, ) from sglang.srt.utils.nvtx_pytorch_hooks import PytHooks from sglang.srt.utils.offloader import ( @@ -199,11 +196,9 @@ from sglang.srt.weight_sync.tensor_bucket import ( FlattenedTensorMetadata, ) -_is_cuda = is_cuda() _is_hip = is_hip() _is_npu = is_npu() _is_cpu_amx_available = cpu_has_amx_support() -_is_xpu_xmx_available = xpu_has_xmx_support() if _is_npu: from sglang.srt.hardware_backend.npu.utils import init_npu_backend @@ -372,8 +367,9 @@ class ModelRunner: self._weight_checker = WeightChecker(model_runner=self) - if get_bool_env_var("SGLANG_DETECT_SLOW_RANK"): + if envs.SGLANG_DETECT_SLOW_RANK.get(): slow_rank_detector.execute() + # Init mindspore running environment when model impl is "mindspore" self.init_mindspore_runner() @@ -431,9 +427,7 @@ class ModelRunner: moe_ep_rank=self.moe_ep_rank, ) ) - if self.tp_rank == 0 and get_bool_env_var( - "SGLANG_LOG_EXPERT_LOCATION_METADATA" - ): + if self.tp_rank == 0 and envs.SGLANG_LOG_EXPERT_LOCATION_METADATA.get(): logger.info( f"Initial expert_location_metadata: {get_global_expert_location_metadata()}" ) @@ -556,12 +550,11 @@ class ModelRunner: enable_batch_invariant_mode() + # Deduce KV cache dtype + self.configure_kv_cache_dtype() + # Init memory pool and attention backends - self.init_memory_pool( - min_per_gpu_memory, - server_args.max_running_requests, - server_args.max_total_tokens, - ) + self.init_memory_pool(min_per_gpu_memory, server_args) # Init max running requests self.max_running_requests = min( @@ -834,16 +827,12 @@ class ModelRunner: local_gpu_memory = get_available_gpu_memory(self.device, self.gpu_id) if self.tp_size > 1 and not self.is_draft_worker: if min_per_gpu_memory < local_gpu_memory * 0.9: - if get_bool_env_var("SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK"): - logger.warning( - "The memory capacity is unbalanced. Some GPUs may be occupied by other processes. " - f"{min_per_gpu_memory=}, {local_gpu_memory=}, {local_gpu_memory * 0.9=}" - ) + msg = "The memory capacity is unbalanced. Some GPUs may be occupied by other processes. " + msg += f"{min_per_gpu_memory=}, {local_gpu_memory=}, {local_gpu_memory * 0.9=}" + if envs.SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK.get(): + raise RuntimeError(msg) else: - raise ValueError( - "The memory capacity is unbalanced. Some GPUs may be occupied by other processes. " - f"{min_per_gpu_memory=}, {local_gpu_memory=}, {local_gpu_memory * 0.9=}" - ) + logger.warning(msg) logger.info( f"Init torch distributed ends. mem usage={(before_avail_memory - local_gpu_memory):.2f} GB" @@ -1500,30 +1489,13 @@ class ModelRunner: return result - def profile_max_num_token(self, total_gpu_memory: int): - available_gpu_memory = get_available_gpu_memory( - self.device, - self.gpu_id, - distributed=get_world_group().world_size > 1, - cpu_group=get_world_group().cpu_group, - ) - if self.is_draft_worker: - num_layers = getattr( - self.model_config.hf_config, - "num_nextn_predict_layers", - self.num_effective_layers, - ) - elif config := self.mambaish_config: - num_layers = len(config.full_attention_layer_ids) - elif self.model_config.full_attention_layer_ids: - num_layers = len(self.model_config.full_attention_layer_ids) - else: - num_layers = self.num_effective_layers + def get_cell_size_per_token(self, num_layers: int) -> int: + kv_size = torch._utils._element_size(self.kv_cache_dtype) if self.use_mla_backend: cell_size = ( (self.model_config.kv_lora_rank + self.model_config.qk_rope_head_dim) * num_layers - * torch._utils._element_size(self.kv_cache_dtype) + * kv_size ) if is_float4_e2m1fn_x2(self.kv_cache_dtype): # kv_scale_buffer @@ -1537,7 +1509,7 @@ class ModelRunner: // scale_block_size ) * num_layers - * torch._utils._element_size(self.kv_cache_dtype) + * kv_size ) # Add indexer KV cache overhead for NSA models (DeepSeek V3.2) @@ -1556,7 +1528,7 @@ class ModelRunner: self.model_config.get_num_kv_heads(get_attention_tp_size()) * (self.model_config.head_dim + self.model_config.v_head_dim) * num_layers - * torch._utils._element_size(self.kv_cache_dtype) + * kv_size ) if is_float4_e2m1fn_x2(self.kv_cache_dtype): @@ -1566,14 +1538,7 @@ class ModelRunner: n = self.model_config.get_num_kv_heads(get_attention_tp_size()) k = self.model_config.head_dim cell_size = (cell_size // 2) + ( - ( - n - * k - * num_layers - * 2 - * torch._utils._element_size(self.kv_cache_dtype) - ) - // scale_block_size + (n * k * num_layers * 2 * kv_size) // scale_block_size ) if self.model_config.hf_config.architectures[0] == "MiMoV2FlashForCausalLM": @@ -1584,8 +1549,34 @@ class ModelRunner: + self.model_config.hf_text_config.swa_v_head_dim ) * len(self.model_config.swa_attention_layer_ids) - * torch._utils._element_size(self.kv_cache_dtype) + * kv_size ) + return cell_size + + def profile_max_num_token(self, total_gpu_memory: int): + available_gpu_memory = get_available_gpu_memory( + self.device, + self.gpu_id, + distributed=get_world_group().world_size > 1, + cpu_group=get_world_group().cpu_group, + ) + + # Get the number of layers used for KV cache calculation + if self.is_draft_worker: + num_layers = getattr( + self.model_config.hf_config, + "num_nextn_predict_layers", + self.num_effective_layers, + ) + elif mambaish := self.mambaish_config: + num_layers = len(mambaish.full_attention_layer_ids) + elif self.model_config.full_attention_layer_ids: + num_layers = len(self.model_config.full_attention_layer_ids) + else: + num_layers = self.num_effective_layers + + cell_size = self.get_cell_size_per_token(num_layers) + rest_memory = available_gpu_memory - total_gpu_memory * ( 1 - self.mem_fraction_static ) @@ -1787,13 +1778,7 @@ class ModelRunner: return False return True - def init_memory_pool( - self, - total_gpu_memory: int, - max_num_reqs: Optional[int] = None, - max_total_tokens: Optional[int] = None, - ): - # Determine the kv cache dtype + def configure_kv_cache_dtype(self): if self.server_args.kv_cache_dtype == "auto": quant_config = getattr(self.model, "quant_config", None) kv_cache_quant_algo = getattr(quant_config, "kv_cache_quant_algo", None) @@ -1835,14 +1820,11 @@ class ModelRunner: log_info_on_rank0(logger, f"Using KV cache dtype: {self.kv_cache_dtype}") + def init_memory_pool(self, total_gpu_memory: int, server_args: ServerArgs): + max_num_reqs = server_args.max_running_requests + max_total_tokens = server_args.max_total_tokens self.max_total_num_tokens = self.profile_max_num_token(total_gpu_memory) - if (small_kv_size := envs.SGLANG_CI_SMALL_KV_SIZE.get()) > 0: - logger.info( - f"Use a small KV cache pool size ({small_kv_size}) for local tests" - ) - self.max_total_num_tokens = small_kv_size - if max_num_reqs is None: max_num_reqs = min( max( diff --git a/test/registered/spec/eagle/test_eagle_infer_b.py b/test/registered/spec/eagle/test_eagle_infer_b.py index 48ae2c96d..fb5c08628 100644 --- a/test/registered/spec/eagle/test_eagle_infer_b.py +++ b/test/registered/spec/eagle/test_eagle_infer_b.py @@ -290,14 +290,16 @@ class TestEAGLEServerBasic(EagleServerBase): class TestEAGLERetract(TestEAGLEServerBasic): - extra_args = ["--chunked-prefill-size", 128, "--max-running-requests", 64] + extra_args = [ + "--chunked-prefill-size=128", + "--max-running-requests=64", + "--max-total-tokens=4500", # Set a smaller KV cache to trigger retract more easily + ] @classmethod def setUpClass(cls): # These config helps find a leak. - with envs.SGLANG_TEST_RETRACT.override( - True - ), envs.SGLANG_CI_SMALL_KV_SIZE.override(4500): + with envs.SGLANG_TEST_RETRACT.override(True): super().setUpClass()