chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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

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(