[Auto Sync] Update weight_utils.py (20260212) (#18692)

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Dan Zheng <dzheng@x.ai>

python/sglang/srt/layers/logits_processor.py

@@ -55,7 +55,7 @@ from sglang.srt.model_executor.forward_batch_info import (
     ForwardMode,
 )
 from sglang.srt.server_args import get_global_server_args
-from sglang.srt.utils import is_npu, use_intel_amx_backend
+from sglang.srt.utils.common import is_npu, use_intel_amx_backend
 
 logger = logging.getLogger(__name__)
 
@@ -89,8 +89,8 @@ class LogitsProcessorOutput:
     # The logprobs of input tokens.        shape: [#token]
     input_token_logprobs: Optional[torch.Tensor] = None
     # The logprobs and ids of the top-k tokens in input positions.  shape: [#seq, #token, k]
-    input_top_logprobs_val: List = None
-    input_top_logprobs_idx: List = None
+    input_top_logprobs_val: Optional[List] = None
+    input_top_logprobs_idx: Optional[List] = None
     # The logprobs and ids of the requested token ids in input positions. shape: [#seq, n] (n is the number of requested token ids)
     # Can contain either lists or GPU tensors (for delayed GPU-to-CPU transfer optimization)
     input_token_ids_logprobs_val: Optional[List[Union[List[float], torch.Tensor]]] = (

python/sglang/srt/model_loader/weight_utils.py

@@ -751,8 +751,12 @@ def safetensors_weights_iterator(
         if disable_mmap:
             with open(st_file, "rb") as f:
                 result = safetensors.torch.load(f.read())
-                for name, param in result.items():
-                    yield name, param
+                # NOTE(xiuyu): safetensors.torch.load() may return keys in a
+                # different order per TP-shard file. Sort to ensure that cross-rank
+                # collectives (e.g. all_gather in sync_quantize_weight) stay
+                # aligned across TP ranks.
+                for name in sorted(result.keys()):
+                    yield name, result[name]
         else:
             with safetensors.safe_open(st_file, framework="pt", device="cpu") as f:
                 for name in f.keys():
@@ -846,23 +850,31 @@ def multi_thread_safetensors_weights_iterator(
         return safetensors.torch.load_file(st_file, device="cpu")
 
     with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
-        futures = [executor.submit(_load_file, st_file) for st_file in hf_weights_files]
+        # Yield in deterministic order to keep cross-rank collectives aligned.
+        future_to_idx = {
+            executor.submit(_load_file, st_file): idx
+            for idx, st_file in enumerate(hf_weights_files)
+        }
+        results_by_idx: Dict[int, dict] = {}
 
         if enable_tqdm:
             futures_iter = tqdm(
-                concurrent.futures.as_completed(futures),
+                concurrent.futures.as_completed(future_to_idx),
                 total=len(hf_weights_files),
                 desc="Multi-thread loading shards",
                 disable=not enable_tqdm,
                 bar_format=BAR_FORMAT,
             )
         else:
-            futures_iter = concurrent.futures.as_completed(futures)
+            futures_iter = concurrent.futures.as_completed(future_to_idx)
 
         for future in futures_iter:
-            state_dict = future.result()
-            for name, param in state_dict.items():
-                yield name, param
+            results_by_idx[future_to_idx[future]] = future.result()
+
+        for idx in range(len(hf_weights_files)):
+            state_dict = results_by_idx[idx]
+            for name in sorted(state_dict.keys()):
+                yield name, state_dict[name]
 
 
 def _load_pt_file(bin_file: str) -> dict: