minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[Core] Clarify memory variable naming in model runner (#20060)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-03-06 14:00:46 -08:00
committed by GitHub
parent ac453b253f
commit 604db4471d
2 changed files with 14 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_executor/model_runner.py
+9 -9
View File
@@ -391,8 +391,8 @@ class ModelRunner(ModelRunnerKVCacheMixin):
# Initialize MooncakeTransferEngine # Initialize MooncakeTransferEngine
self.init_shared_mooncake_transfer_engine() self.init_shared_mooncake_transfer_engine()
# Get memory before model loading # Get available memory before model loading
min_per_gpu_memory = self.init_torch_distributed() pre_model_load_memory = self.init_torch_distributed()
# Init forward stream for overlap schedule # Init forward stream for overlap schedule
self.forward_stream = torch.get_device_module(self.device).Stream() self.forward_stream = torch.get_device_module(self.device).Stream()
@@ -413,7 +413,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
deep_gemm_wrapper.update_deep_gemm_config(gpu_id, server_args) deep_gemm_wrapper.update_deep_gemm_config(gpu_id, server_args)
# Initialize the model runner # Initialize the model runner
self.initialize(min_per_gpu_memory) self.initialize(pre_model_load_memory)
self.check_quantized_moe_compatibility() self.check_quantized_moe_compatibility()
if self.is_multimodal: if self.is_multimodal:
@@ -447,7 +447,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
port=self.dist_port, port=self.dist_port,
) )
def initialize(self, min_per_gpu_memory: float): def initialize(self, pre_model_load_memory: float):
server_args = self.server_args server_args = self.server_args
self.memory_saver_adapter = TorchMemorySaverAdapter.create( self.memory_saver_adapter = TorchMemorySaverAdapter.create(
@@ -600,7 +600,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
self.configure_kv_cache_dtype() self.configure_kv_cache_dtype()
# Init memory pool and attention backends # Init memory pool and attention backends
self.init_memory_pool(min_per_gpu_memory) self.init_memory_pool(pre_model_load_memory)
# Init max running requests # Init max running requests
self.max_running_requests = min( self.max_running_requests = min(
@@ -826,7 +826,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
if is_npu(): if is_npu():
register_sgl_tp_rank(self.gpu_id) register_sgl_tp_rank(self.gpu_id)
min_per_gpu_memory = get_available_gpu_memory( pre_model_load_memory = get_available_gpu_memory(
self.device, self.device,
self.gpu_id, self.gpu_id,
distributed=get_world_group().world_size > 1, distributed=get_world_group().world_size > 1,
@@ -839,9 +839,9 @@ class ModelRunner(ModelRunnerKVCacheMixin):
# Check memory for tensor parallelism # Check memory for tensor parallelism
local_gpu_memory = get_available_gpu_memory(self.device, self.gpu_id) local_gpu_memory = get_available_gpu_memory(self.device, self.gpu_id)
if self.tp_size > 1 and not self.is_draft_worker: if self.tp_size > 1 and not self.is_draft_worker:
if min_per_gpu_memory < local_gpu_memory * 0.9: if pre_model_load_memory < local_gpu_memory * 0.9:
msg = "The memory capacity is unbalanced. Some GPUs may be occupied by other processes. " 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=}" msg += f"{pre_model_load_memory=}, {local_gpu_memory=}, {local_gpu_memory * 0.9=}"
if envs.SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK.get(): if envs.SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK.get():
raise RuntimeError(msg) raise RuntimeError(msg)
else: else:
@@ -851,7 +851,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
f"Init torch distributed ends. elapsed={time.perf_counter() - tic:.2f} s, " f"Init torch distributed ends. elapsed={time.perf_counter() - tic:.2f} s, "
f"mem usage={(before_avail_memory - local_gpu_memory):.2f} GB" f"mem usage={(before_avail_memory - local_gpu_memory):.2f} GB"
) )
return min_per_gpu_memory return pre_model_load_memory
def init_shared_mooncake_transfer_engine(self): def init_shared_mooncake_transfer_engine(self):
""" """
python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py
+5 -5
View File
@@ -113,8 +113,8 @@ class ModelRunnerKVCacheMixin:
) )
return cell_size return cell_size
def profile_max_num_token(self: ModelRunner, total_gpu_memory: int): def profile_max_num_token(self: ModelRunner, pre_model_load_memory: int):
available_gpu_memory = get_available_gpu_memory( post_model_load_memory = get_available_gpu_memory(
self.device, self.device,
self.gpu_id, self.gpu_id,
distributed=get_world_group().world_size > 1, distributed=get_world_group().world_size > 1,
@@ -140,7 +140,7 @@ class ModelRunnerKVCacheMixin:
cell_size = self.get_cell_size_per_token(num_layers) cell_size = self.get_cell_size_per_token(num_layers)
rest_memory = available_gpu_memory - total_gpu_memory * ( rest_memory = post_model_load_memory - pre_model_load_memory * (
1 - self.mem_fraction_static 1 - self.mem_fraction_static
) )
if self.mambaish_config is not None: if self.mambaish_config is not None:
@@ -675,10 +675,10 @@ class ModelRunnerKVCacheMixin:
self.token_to_kv_pool_allocator.full_to_swa_index_mapping self.token_to_kv_pool_allocator.full_to_swa_index_mapping
) )
def init_memory_pool(self: ModelRunner, total_gpu_memory: int): def init_memory_pool(self: ModelRunner, pre_model_load_memory: int):
max_num_reqs = self.server_args.max_running_requests max_num_reqs = self.server_args.max_running_requests
max_total_tokens = self.server_args.max_total_tokens max_total_tokens = self.server_args.max_total_tokens
self.max_total_num_tokens = self.profile_max_num_token(total_gpu_memory) self.max_total_num_tokens = self.profile_max_num_token(pre_model_load_memory)
if max_num_reqs is None: if max_num_reqs is None:
max_num_reqs = min( max_num_reqs = min(