optm(checkpoint-engine): disable multi-thread loading when update weights (#12374)

Signed-off-by: Yang Kaiyong <yangkaiyong.yky@antgroup.com>

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)