Cleanup ModelRunner (#15802)

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
 

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_"):

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(
+            if self.tp_rank == 0 and envs.SGLANG_LOG_EXPERT_LOCATION_METADATA.get():
-                "SGLANG_LOG_EXPERT_LOCATION_METADATA"
-            ):
                 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(
+        self.init_memory_pool(min_per_gpu_memory, server_args)
-            min_per_gpu_memory,
-            server_args.max_running_requests,
-            server_args.max_total_tokens,
-        )
 
         # 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"):
+                msg = "The memory capacity is unbalanced. Some GPUs may be occupied by other processes. "
-                    logger.warning(
+                msg += f"{min_per_gpu_memory=}, {local_gpu_memory=}, {local_gpu_memory * 0.9=}"
-                        "The memory capacity is unbalanced. Some GPUs may be occupied by other processes. "
+                if envs.SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK.get():
-                        f"{min_per_gpu_memory=}, {local_gpu_memory=}, {local_gpu_memory * 0.9=}"
+                    raise RuntimeError(msg)
-                    )
                 else:
-                    raise ValueError(
+                    logger.warning(msg)
-                        "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.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):
+    def get_cell_size_per_token(self, num_layers: int) -> int:
-        available_gpu_memory = get_available_gpu_memory(
+        kv_size = torch._utils._element_size(self.kv_cache_dtype)
-            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
         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 * kv_size) // scale_block_size
-                        n
-                        * k
-                        * num_layers
-                        * 2
-                        * torch._utils._element_size(self.kv_cache_dtype)
-                    )
-                    // 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(
+    def configure_kv_cache_dtype(self):
-        self,
-        total_gpu_memory: int,
-        max_num_reqs: Optional[int] = None,
-        max_total_tokens: Optional[int] = None,
-    ):
-        # Determine the kv cache dtype
         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(

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(
+        with envs.SGLANG_TEST_RETRACT.override(True):
-            True
-        ), envs.SGLANG_CI_SMALL_KV_SIZE.override(4500):
             super().setUpClass()