diff --git a/python/sglang/srt/model_loader/utils.py b/python/sglang/srt/model_loader/utils.py index 3e315e10a..2511c61cc 100644 --- a/python/sglang/srt/model_loader/utils.py +++ b/python/sglang/srt/model_loader/utils.py @@ -1,9 +1,10 @@ # Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/model_executor/model_loader/utils.py """Utilities for selecting and loading models.""" +import concurrent.futures import contextlib import logging -from typing import Tuple, Type +from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Type import torch import transformers @@ -116,3 +117,43 @@ def post_load_weights(model: nn.Module, model_config: ModelConfig): model.post_load_weights(is_nextn=True) else: model.post_load_weights() + + +def should_async_load(weight: torch.Tensor) -> bool: + """Return True if we should load the given weight asynchronously. + + For host (CPU) tensors, using a threadpool can overlap H2D copies + and improve throughput. For device tensors, threading often adds overhead + (e.g., GIL contention) without benefit, so we do it synchronously. + """ + device = getattr(weight, "device", None) + if device is None: + return False + return device.type == "cpu" + + +def maybe_executor_submit( + *, + executor: concurrent.futures.ThreadPoolExecutor, + futures: List[concurrent.futures.Future], + use_async: bool, + func: Callable[..., Any], + func_args: Iterable[Any] = (), + func_kwargs: Optional[Dict[str, Any]] = None, +) -> None: + """Submit a task to the executor if async loading is enabled. + + Parameters (keyword-only): + - executor: ThreadPoolExecutor used to submit background tasks + - futures: a list collecting the submitted Future objects + - use_async: whether to submit to executor or run inline + - func: the callable to run + - func_args: positional args for the callable (defaults to empty tuple) + - func_kwargs: keyword args for the callable (defaults to empty dict) + """ + if func_kwargs is None: + func_kwargs = {} + if use_async: + futures.append(executor.submit(func, *func_args, **func_kwargs)) + else: + func(*func_args, **func_kwargs) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index d9c4eca4d..28dd58483 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -112,6 +112,7 @@ from sglang.srt.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors +from sglang.srt.model_loader.utils import maybe_executor_submit, should_async_load from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.server_args import get_global_server_args from sglang.srt.single_batch_overlap import SboFlags @@ -3505,6 +3506,7 @@ class DeepseekV2ForCausalLM(nn.Module): params_dict = dict(self.named_parameters()) weight_names = [] for name, loaded_weight in weights: + use_async_loading = should_async_load(loaded_weight) layer_id = get_layer_id(name) if ( layer_id is not None @@ -3572,8 +3574,12 @@ class DeepseekV2ForCausalLM(nn.Module): continue param = params_dict[name] weight_loader = param.weight_loader - futures.append( - executor.submit(weight_loader, param, loaded_weight, shard_id) + maybe_executor_submit( + executor=executor, + futures=futures, + use_async=use_async_loading, + func=weight_loader, + func_args=(param, loaded_weight, shard_id), ) break else: @@ -3584,15 +3590,20 @@ class DeepseekV2ForCausalLM(nn.Module): name = name.replace(weight_name, param_name) param = params_dict[name] weight_loader = param.weight_loader - futures.append( - executor.submit( - weight_loader, + maybe_executor_submit( + executor=executor, + futures=futures, + use_async=use_async_loading, + func=weight_loader, + func_args=( param, loaded_weight, name, - shard_id=shard_id, - expert_id=expert_id, - ) + ), + func_kwargs={ + "shard_id": shard_id, + "expert_id": expert_id, + }, ) break else: @@ -3652,8 +3663,12 @@ class DeepseekV2ForCausalLM(nn.Module): weight_loader = getattr( param, "weight_loader", default_weight_loader ) - futures.append( - executor.submit(weight_loader, param, fused_weight) + maybe_executor_submit( + executor=executor, + futures=futures, + use_async=use_async_loading, + func=weight_loader, + func_args=(param, fused_weight), ) cached_a_proj.pop(q_a_proj_name) cached_a_proj.pop(kv_a_proj_name) @@ -3699,8 +3714,13 @@ class DeepseekV2ForCausalLM(nn.Module): weight_loader = getattr( param, "weight_loader", default_weight_loader ) - futures.append( - executor.submit(weight_loader, param, fused_weight) + maybe_executor_submit( + executor, + futures, + use_async_loading, + weight_loader, + param, + fused_weight, ) cached_wk_and_weights_proj.pop(wk_name) cached_wk_and_weights_proj.pop(weights_proj_name) @@ -3725,8 +3745,12 @@ class DeepseekV2ForCausalLM(nn.Module): weight_loader = getattr( param, "weight_loader", default_weight_loader ) - futures.append( - executor.submit(weight_loader, param, loaded_weight) + maybe_executor_submit( + executor=executor, + futures=futures, + use_async=use_async_loading, + func=weight_loader, + func_args=(param, loaded_weight), ) # Wait for all tasks to complete and raise any exceptions. diff --git a/python/sglang/srt/models/longcat_flash.py b/python/sglang/srt/models/longcat_flash.py index 23bfed025..84aeb8b30 100644 --- a/python/sglang/srt/models/longcat_flash.py +++ b/python/sglang/srt/models/longcat_flash.py @@ -80,6 +80,7 @@ from sglang.srt.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.utils import maybe_executor_submit, should_async_load from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA from sglang.srt.server_args import get_global_server_args @@ -854,6 +855,7 @@ class LongcatFlashForCausalLM(nn.Module): params_dict = dict(self.named_parameters()) weight_names = [] for name, loaded_weight in weights: + use_async_loading = should_async_load(loaded_weight) if "mtp" in name: continue weight_names.append(name) @@ -877,8 +879,12 @@ class LongcatFlashForCausalLM(nn.Module): continue param = params_dict[name] weight_loader = param.weight_loader - futures.append( - executor.submit(weight_loader, param, loaded_weight, shard_id) + maybe_executor_submit( + executor=executor, + futures=futures, + use_async=use_async_loading, + func=weight_loader, + func_args=(param, loaded_weight, shard_id), ) break else: @@ -889,15 +895,16 @@ class LongcatFlashForCausalLM(nn.Module): name = name.replace(weight_name, param_name) param = params_dict[name] weight_loader = param.weight_loader - futures.append( - executor.submit( - weight_loader, - param, - loaded_weight, - name, - shard_id=shard_id, - expert_id=expert_id, - ) + maybe_executor_submit( + executor=executor, + futures=futures, + use_async=use_async_loading, + func=weight_loader, + func_args=(param, loaded_weight, name), + func_kwargs={ + "shard_id": shard_id, + "expert_id": expert_id, + }, ) break else: @@ -951,8 +958,12 @@ class LongcatFlashForCausalLM(nn.Module): weight_loader = getattr( param, "weight_loader", default_weight_loader ) - futures.append( - executor.submit(weight_loader, param, fused_weight) + maybe_executor_submit( + executor=executor, + futures=futures, + use_async=use_async_loading, + func=weight_loader, + func_args=(param, fused_weight), ) cached_a_proj.pop(q_a_proj_name) cached_a_proj.pop(kv_a_proj_name) @@ -977,8 +988,12 @@ class LongcatFlashForCausalLM(nn.Module): weight_loader = getattr( param, "weight_loader", default_weight_loader ) - futures.append( - executor.submit(weight_loader, param, loaded_weight) + maybe_executor_submit( + executor=executor, + futures=futures, + use_async=use_async_loading, + func=weight_loader, + func_args=(param, loaded_weight), ) # Wait for all tasks to complete and raise any exceptions.