chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Cleanup ModelRunner (#15802)

Browse Source
This commit is contained in:
Liangsheng Yin
2025-12-25 18:13:30 +08:00
committed by GitHub
parent e6ce16a4c2
commit f4e835af2f
4 changed files with 62 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
docs/references/environment_variables.md
View File

@@ -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

7
python/sglang/srt/environ.py
View File

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

116
python/sglang/srt/model_executor/model_runner.py
View File

@@ -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(

10
test/registered/spec/eagle/test_eagle_infer_b.py
View File

@@ -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()