Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
hnyls2002 merged 2 commits into from
Mar 6, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
21f945b
clarify the memory calculation in model runner
hnyls2002 Mar 6, 2026
39ff3a9
Merge branch 'main' into lsyin/tiny-cleanup-num-reqs
hnyls2002 Mar 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 9 additions & 9 deletions python/sglang/srt/model_executor/model_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -391,8 +391,8 @@ def __init__(
# Initialize MooncakeTransferEngine
self.init_shared_mooncake_transfer_engine()

# Get memory before model loading
min_per_gpu_memory = self.init_torch_distributed()
# Get available memory before model loading
pre_model_load_memory = self.init_torch_distributed()

# Init forward stream for overlap schedule
self.forward_stream = torch.get_device_module(self.device).Stream()
Expand All @@ -413,7 +413,7 @@ def __init__(
deep_gemm_wrapper.update_deep_gemm_config(gpu_id, server_args)

# Initialize the model runner
self.initialize(min_per_gpu_memory)
self.initialize(pre_model_load_memory)
self.check_quantized_moe_compatibility()

if self.is_multimodal:
Expand Down Expand Up @@ -447,7 +447,7 @@ def init_mindspore_runner(self):
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

self.memory_saver_adapter = TorchMemorySaverAdapter.create(
Expand Down Expand Up @@ -600,7 +600,7 @@ def initialize(self, min_per_gpu_memory: float):
self.configure_kv_cache_dtype()

# 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
self.max_running_requests = min(
Expand Down Expand Up @@ -826,7 +826,7 @@ def _(data, dim):
if is_npu():
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.gpu_id,
distributed=get_world_group().world_size > 1,
Expand All @@ -839,9 +839,9 @@ def _(data, dim):
# Check memory for tensor parallelism
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 pre_model_load_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=}"
msg += f"{pre_model_load_memory=}, {local_gpu_memory=}, {local_gpu_memory * 0.9=}"
if envs.SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK.get():
raise RuntimeError(msg)
else:
Expand All @@ -851,7 +851,7 @@ def _(data, dim):
f"Init torch distributed ends. elapsed={time.perf_counter() - tic:.2f} s, "
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):
"""
Expand Down
10 changes: 5 additions & 5 deletions python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -113,8 +113,8 @@ def get_cell_size_per_token(self: ModelRunner, num_layers: int) -> int:
)
return cell_size

def profile_max_num_token(self: ModelRunner, total_gpu_memory: int):
available_gpu_memory = get_available_gpu_memory(
def profile_max_num_token(self: ModelRunner, pre_model_load_memory: int):
post_model_load_memory = get_available_gpu_memory(
self.device,
self.gpu_id,
distributed=get_world_group().world_size > 1,
Expand All @@ -140,7 +140,7 @@ def profile_max_num_token(self: ModelRunner, total_gpu_memory: int):

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
)
if self.mambaish_config is not None:
Expand Down Expand Up @@ -675,10 +675,10 @@ def _init_pools(self: ModelRunner, max_num_reqs: int):
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_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:
max_num_reqs = min(
Expand Down
Loading