[core/loader] Add presharded load format (#24256)

Co-authored-by: Shu Wang <shuwanguc@google.com>

python/pyproject.toml

@@ -87,6 +87,7 @@ dependencies = [
   "uvloop",
   "watchfiles",
   "xgrammar==0.2.1",
+  "xxhash",
   "zstandard",
 ]
 

python/pyproject_cpu.toml

@@ -69,6 +69,7 @@ dependencies = [
   "triton==3.7.0",
   "uvicorn",
   "uvloop",
+  "xxhash",
   "xgrammar==0.2.1",
   "zstandard",
 ]

python/pyproject_npu.toml

@@ -65,6 +65,7 @@ dependencies = [
   "transformers==5.12.1",
   "uvicorn",
   "uvloop",
+  "xxhash",
   "xgrammar==0.2.1",
 ]
 

python/pyproject_other.toml

@@ -66,6 +66,7 @@ runtime_common = [
   "transformers==5.12.1",
   "uvicorn",
   "uvloop",
+  "xxhash",
   "xgrammar==0.2.1",
 ]
 

python/pyproject_xpu.toml

@@ -67,6 +67,7 @@ dependencies = [
   "tqdm",
   "transformers==5.12.1",
   "uvicorn",
+  "xxhash",
   "uvloop",
   # "xgrammar==0.2.1", xgrammar depends on CUDA PyTorch and Triton only
 ]

python/sglang/srt/configs/load_config.py

@@ -21,6 +21,7 @@ class LoadFormat(str, enum.Enum):
     NPCACHE = "npcache"
     DUMMY = "dummy"
     SHARDED_STATE = "sharded_state"
+    PRESHARDED = "presharded"
     GGUF = "gguf"
     BITSANDBYTES = "bitsandbytes"
     MISTRAL = "mistral"

python/sglang/srt/model_loader/loader.py

@@ -6,15 +6,18 @@ from __future__ import annotations
 
 # ruff: noqa: SIM117
 import collections
+import concurrent.futures
 import dataclasses
 import fnmatch
 import gc
 import glob
+import hashlib
 import json
 import logging
 import math
 import os
 import re
+import shutil
 import socket
 import tempfile
 import threading
@@ -1497,6 +1500,15 @@ class ShardedStateLoader(BaseModelLoader):
         result: Dict[str, torch.Tensor] = {}
         for group in same_storage_groups.values():
             for k, t in group:
+                if not t.is_contiguous():
+                    # End-pointer dedup assumes a flat view; non-contiguous
+                    # tensors (e.g. produced by
+                    # ``.transpose(...).contiguous().transpose(...)`` in some
+                    # quant ``post_load_weights`` paths) cannot be flattened
+                    # via ``view(-1)``. Include them directly; downstream
+                    # writers call ``.contiguous()`` before save.
+                    result[k] = t
+                    continue
                 a, b = t.data_ptr(), get_end_ptr(t)
                 for k2, t2 in group:
                     if not t2.is_contiguous():
@@ -1629,6 +1641,839 @@ class ShardedStateLoader(BaseModelLoader):
             )
 
 
+class PreshardedModelLoader(DefaultModelLoader):
+    """Dump/reload post-process weights under ``<model_path>/presharded/<subdir>/``.
+
+    Optional roots in ``model_loader_extra_config`` (subdir still appended):
+    ``presharded_path`` (target), ``draft_presharded_path`` (speculative draft).
+    Dump dir must be shared across ranks/nodes.
+    """
+
+    DEFAULT_SUBDIR = "presharded"
+    MAX_FILE_BYTES = 20 * (1024**3)
+    CHECKSUM_FILENAME = "checksum.json"
+    READY_FILENAME = "READY"
+    TMP_SUBDIR = "_tmp_presharding"
+    PLAN_VERSION = 1
+    DEFAULT_HASH_NUM_THREADS = 8
+    _CONTENT_HASH_HEX_LEN = 32
+
+    def __init__(self, load_config: LoadConfig):
+        extra = (
+            {}
+            if load_config.model_loader_extra_config is None
+            else dict(load_config.model_loader_extra_config)
+        )
+        self._presharded_path_override = extra.pop("presharded_path", None)
+        self._draft_presharded_path_override = extra.pop("draft_presharded_path", None)
+        self._max_file_bytes = int(extra.pop("max_file_bytes", self.MAX_FILE_BYTES))
+        self._hash_num_threads = int(
+            extra.pop("hash_num_threads", self.DEFAULT_HASH_NUM_THREADS)
+        )
+        self._verify_on_load = bool(extra.pop("verify_on_load", False))
+        load_config.model_loader_extra_config = extra
+        load_config.load_format = LoadFormat.AUTO
+        super().__init__(load_config)
+
+    def download_model(self, model_config: ModelConfig) -> None:
+        presharded_dir = self._presharded_dir(model_config)
+        if not self._presharded_ready(presharded_dir):
+            super().download_model(model_config)
+
+    def load_model(
+        self,
+        *,
+        model_config: ModelConfig,
+        device_config: DeviceConfig,
+    ) -> nn.Module:
+        shard_config = self._collect_shard_config(model_config)
+        presharded_dir = self._presharded_dir(model_config, shard_config)
+        if self._presharded_ready(presharded_dir) and self._shard_config_matches(
+            presharded_dir, shard_config
+        ):
+            logger.info("Loading from presharded checkpoint at %s", presharded_dir)
+            return self._load_from_presharded(
+                model_config, device_config, presharded_dir
+            )
+        logger.info(
+            "No presharded checkpoint at %s; doing first-time load and dump.",
+            presharded_dir,
+        )
+        return self._first_time_load_and_dump(
+            model_config, device_config, presharded_dir, shard_config
+        )
+
+    @classmethod
+    def _presharded_ready(cls, presharded_dir: str) -> bool:
+        return os.path.isfile(os.path.join(presharded_dir, cls.READY_FILENAME))
+
+    def _presharded_dir(
+        self,
+        model_config: ModelConfig,
+        shard_config: Optional[Dict[str, Any]] = None,
+    ) -> str:
+        if shard_config is None:
+            shard_config = self._collect_shard_config(model_config)
+        subfolder = self._build_subfolder_name(shard_config)
+        if model_config.is_draft_model:
+            root = self._draft_presharded_path_override
+        else:
+            root = self._presharded_path_override
+        if root is None:
+            root = os.path.join(model_config.model_path, self.DEFAULT_SUBDIR)
+        return os.path.join(root, subfolder)
+
+    def _collect_shard_config(self, model_config: ModelConfig) -> Dict[str, Any]:
+        def _safe(fn) -> int:
+            try:
+                return fn()
+            except (AssertionError, AttributeError, RuntimeError):
+                return 1
+
+        parallel = get_parallel()
+        server_args = get_server_args()
+        return {
+            "tp": _safe(lambda: parallel.tp_size),
+            "dp": _safe(lambda: parallel.moe_dp_size),
+            "ep": _safe(lambda: parallel.moe_ep_size),
+            "pp": _safe(lambda: parallel.pp_size),
+            "moe_dense_tp_size": server_args.moe_dense_tp_size,
+            "moe_dp_size": server_args.moe_dp_size,
+            "enable_dp_lm_head": server_args.enable_dp_lm_head,
+            "enable_fp32_lm_head": server_args.enable_fp32_lm_head,
+            "quantization": model_config.quantization,
+            "model_dtype": str(model_config.dtype),
+            "ep_num_redundant_experts": server_args.ep_num_redundant_experts,
+            "enable_eplb": server_args.enable_eplb,
+            "init_expert_location": self._normalize_init_expert_location(
+                server_args.init_expert_location
+            ),
+            "structural_signature": self._compute_structural_signature(model_config),
+        }
+
+    @staticmethod
+    def _normalize_init_expert_location(value: Optional[str]) -> Optional[str]:
+        if value is None or value == "trivial":
+            return value
+        if value.endswith((".json", ".pt")) and os.path.isfile(value):
+            h = hashlib.sha1()
+            with open(value, "rb") as f:
+                for chunk in iter(lambda: f.read(1024 * 1024), b""):
+                    h.update(chunk)
+            return f"file:{os.path.basename(value)}:sha1:{h.hexdigest()[:16]}"
+        return value
+
+    def _build_subfolder_name(self, shard_config: Dict[str, Any]) -> str:
+        combined = hashlib.sha1(
+            json.dumps(shard_config, sort_keys=True).encode()
+        ).hexdigest()[:16]
+        return f"TP-{shard_config['tp']}-sig-{combined}"
+
+    def _shard_config_matches(
+        self, presharded_dir: str, shard_config: Dict[str, Any]
+    ) -> bool:
+        try:
+            with open(os.path.join(presharded_dir, self.CHECKSUM_FILENAME)) as f:
+                stored = json.load(f).get("shard_config")
+        except (OSError, ValueError):
+            stored = None
+        current = json.loads(json.dumps(shard_config))
+        if stored == current:
+            return True
+        logger.warning(
+            "Presharded checkpoint at %s was dumped with a different shard "
+            "config than the current launch (stored=%s, current=%s). "
+            "Treating as a cache miss and re-dumping.",
+            presharded_dir,
+            stored,
+            current,
+        )
+        return False
+
+    def _compute_structural_signature(self, model_config: ModelConfig) -> Optional[str]:
+        local_sig = self._compute_local_structural_signature(model_config)
+        return self._make_rank_invariant_structural_signature(local_sig)
+
+    def _compute_local_structural_signature(
+        self, model_config: ModelConfig
+    ) -> Optional[str]:
+        from sglang.srt.layers.rotary_embedding.factory import _ROPE_DICT
+
+        def _clear_meta_rope_cache() -> None:
+            meta_keys = [
+                k
+                for k, v in _ROPE_DICT.items()
+                if any(p.device.type == "meta" for p in v.parameters())
+                or any(b.device.type == "meta" for b in v.buffers())
+            ]
+            for k in meta_keys:
+                del _ROPE_DICT[k]
+
+        try:
+            quant_config = _get_quantization_config(model_config, self.load_config)
+            with set_default_torch_dtype(model_config.dtype):
+                with torch.device("meta"):
+                    meta_model = _initialize_model(
+                        model_config, self.load_config, quant_config
+                    )
+                state_dict = meta_model.state_dict()
+                sig_input = sorted(
+                    (name, tuple(t.shape), str(t.dtype))
+                    for name, t in state_dict.items()
+                )
+            del meta_model
+            return self._hash_structural_signature(sig_input)
+        except Exception as e:
+            logger.warning(
+                "Failed to build structural signature for presharded cache key "
+                "(model_type=%s): %s",
+                getattr(
+                    getattr(model_config, "hf_config", None), "model_type", "unknown"
+                ),
+                e,
+            )
+            return None
+        finally:
+            _clear_meta_rope_cache()
+
+    @classmethod
+    def _make_rank_invariant_structural_signature(
+        cls, local_sig: Optional[str]
+    ) -> Optional[str]:
+        try:
+            from sglang.srt.distributed import get_world_group
+
+            group = get_world_group()
+            if group.world_size <= 1:
+                return local_sig
+            all_sigs = group.all_gather_object(local_sig)
+        except (AssertionError, AttributeError, RuntimeError):
+            return local_sig
+
+        if all(s is None for s in all_sigs):
+            return None
+        return hashlib.sha1(repr(all_sigs).encode()).hexdigest()[:16]
+
+    @staticmethod
+    def _hash_structural_signature(
+        sig_input: List[Tuple[str, Tuple[int, ...], str]],
+    ) -> str:
+        h = hashlib.sha1(repr(sig_input).encode())
+        return h.hexdigest()[:16]
+
+    @staticmethod
+    def _world_rank_and_size() -> Tuple[int, int]:
+        from sglang.srt.distributed import get_world_group
+
+        try:
+            g = get_world_group()
+            return g.rank_in_group, g.world_size
+        except (AssertionError, AttributeError):
+            return 0, 1
+
+    @staticmethod
+    def _world_barrier() -> None:
+        from sglang.srt.distributed import get_world_group
+
+        try:
+            get_world_group().barrier()
+        except (AssertionError, AttributeError):
+            pass
+
+    @staticmethod
+    def _new_content_hasher():
+        import xxhash
+
+        return xxhash.xxh3_128()
+
+    @staticmethod
+    def _hash_tensor(tensor: torch.Tensor) -> str:
+        # CPU copy so concurrent dump workers cannot race CUDA D2H hashing.
+        t = tensor.detach()
+        prefix = str(tuple(t.shape)).encode() + str(t.dtype).encode()
+        h = PreshardedModelLoader._new_content_hasher()
+        h.update(prefix)
+
+        if t.numel() == 0:
+            return h.hexdigest()
+
+        cpu = t.contiguous().to(device="cpu", copy=True).contiguous()
+        flat_u8 = cpu.reshape(-1).view(torch.uint8)
+        h.update(memoryview(flat_u8.numpy()))
+        return h.hexdigest()
+
+    def _verify_rank_checksum(
+        self,
+        verify_hashes: List[Tuple[str, str]],
+        plan: Dict[str, Any],
+        rank: int,
+        presharded_dir: str,
+    ) -> None:
+        expected = plan.get("rank_checksums", {}).get(str(rank))
+        if expected is None:
+            raise ValueError(
+                f"Plan at {presharded_dir} has no rank_checksums entry for "
+                f"rank {rank}; cannot verify. Set "
+                f"--model-loader-extra-config '{{\"verify_on_load\": false}}' "
+                f"to skip verification, or re-dump the checkpoint."
+            )
+
+        total = 0
+        for name, content_hash in verify_hashes:
+            d = PreshardedModelLoader._fold_name_content_digest(name, content_hash)
+            total = (total + int.from_bytes(d[:8], "big")) & 0xFFFFFFFFFFFFFFFF
+        actual = format(total, "016x")
+
+        if actual != expected:
+            raise ValueError(
+                f"Rank-{rank} checksum mismatch for presharded checkpoint at "
+                f"{presharded_dir}: expected {expected}, got {actual}. The "
+                f"checkpoint files may be corrupted; re-dump or skip "
+                f"verification with --model-loader-extra-config "
+                f"'{{\"verify_on_load\": false}}'."
+            )
+
+    @staticmethod
+    def _fold_name_content_digest(name: str, content_hash: str) -> bytes:
+        h = PreshardedModelLoader._new_content_hasher()
+        h.update((name + ":" + content_hash).encode("utf-8"))
+        return h.digest()
+
+    @staticmethod
+    def _collect_extra_tensors(model: nn.Module) -> Dict[str, torch.Tensor]:
+        seen: set = set()
+        param_storages: set = set()
+        for name, tensor in model.state_dict().items():
+            seen.add(name)
+            if tensor.numel() > 0:
+                param_storages.add((tensor.device, tensor.untyped_storage().data_ptr()))
+        extras: Dict[str, torch.Tensor] = {}
+        for module_name, module in model.named_modules():
+            prefix = f"{module_name}." if module_name else ""
+            for attr_name in list(vars(module).keys()):
+                if attr_name.startswith("_"):
+                    continue
+                try:
+                    val = getattr(module, attr_name)
+                except AttributeError:
+                    continue
+                if isinstance(val, torch.Tensor) and not isinstance(
+                    val, torch.nn.Parameter
+                ):
+                    full_name = f"{prefix}{attr_name}"
+                    if full_name in seen:
+                        continue
+                    if val.numel() > 0:
+                        key = (val.device, val.untyped_storage().data_ptr())
+                        if key in param_storages:
+                            continue
+                    extras[full_name] = val
+        return extras
+
+    @staticmethod
+    def _rebind_parameter_aliases(model: nn.Module) -> None:
+        for _, module in model.named_modules():
+            gemma_w = getattr(module, "gemma_weight", None)
+            weight = getattr(module, "weight", None)
+            if (
+                isinstance(gemma_w, torch.Tensor)
+                and isinstance(weight, torch.nn.Parameter)
+                and gemma_w.shape == weight.shape
+            ):
+                torch.add(weight.data, 1.0, out=gemma_w)
+
+            attn = getattr(module, "attn", None)
+            conv1d = getattr(module, "conv1d", None)
+            if attn is None:
+                continue
+            if hasattr(module, "A_log") and hasattr(attn, "A_log"):
+                attn.A_log = module.A_log
+            if hasattr(module, "dt_bias") and hasattr(attn, "dt_bias"):
+                attn.dt_bias = module.dt_bias
+            if conv1d is None:
+                continue
+            cweight = getattr(conv1d, "weight", None)
+            if cweight is not None and hasattr(attn, "conv_weights"):
+                if cweight.dim() == 3 and cweight.size(1) == 1:
+                    attn.conv_weights = cweight.view(cweight.size(0), cweight.size(2))
+                else:
+                    attn.conv_weights = (
+                        cweight.squeeze() if cweight.dim() > 2 else cweight
+                    )
+            if hasattr(conv1d, "bias") and hasattr(attn, "bias"):
+                attn.bias = conv1d.bias
+
+    def _ensure_presharded_dir_writable(self, presharded_dir: str) -> None:
+        rank, _ = self._world_rank_and_size()
+        try:
+            os.makedirs(presharded_dir, exist_ok=True)
+            if rank == 0:
+                probe = os.path.join(presharded_dir, ".presharded_write_probe")
+                last_err: Optional[OSError] = None
+                for _ in range(5):
+                    try:
+                        with open(probe, "w") as f:
+                            f.write("ok")
+                        os.unlink(probe)
+                        last_err = None
+                        break
+                    except OSError as e:
+                        last_err = e
+                        os.makedirs(presharded_dir, exist_ok=True)
+                        time.sleep(0.05)
+                if last_err is not None:
+                    raise last_err
+        except OSError as e:
+            raise RuntimeError(
+                f"Presharded dump directory is not writable: {presharded_dir}. "
+                "Set model_loader_extra_config "
+                '\'{"presharded_path": "..."}\' (or draft_presharded_path for '
+                "the draft model) to a writable shared filesystem path. "
+                f"Original error: {e}"
+            ) from e
+        self._world_barrier()
+
+    def _first_time_load_and_dump(
+        self,
+        model_config: ModelConfig,
+        device_config: DeviceConfig,
+        presharded_dir: str,
+        shard_config: Dict[str, Any],
+    ) -> nn.Module:
+        self._ensure_presharded_dir_writable(presharded_dir)
+        target_device = torch.device(device_config.device)
+        quant_config = _get_quantization_config(model_config, self.load_config)
+        with set_default_torch_dtype(model_config.dtype):
+            with target_device:
+                model = _initialize_model(model_config, self.load_config, quant_config)
+            self.load_weights_and_postprocess(
+                model,
+                self._get_all_weights(model_config, model),
+                target_device,
+            )
+
+            state_dict = dict(model.state_dict())
+            extras = self._collect_extra_tensors(model)
+            self._dump_state_to_disk(state_dict, extras, presharded_dir, shard_config)
+            del state_dict
+            del extras
+            gc.collect()
+
+        self.counter_after_loading_weights = time.perf_counter()
+        return model.eval()
+
+    def _dump_state_to_disk(
+        self,
+        state_dict: Dict[str, torch.Tensor],
+        extras: Dict[str, torch.Tensor],
+        presharded_dir: str,
+        shard_config: Dict[str, Any],
+    ) -> None:
+        rank, world_size = self._world_rank_and_size()
+        tmp_dir = os.path.join(presharded_dir, self.TMP_SUBDIR)
+        if rank == 0:
+            ready_path = os.path.join(presharded_dir, self.READY_FILENAME)
+            if os.path.isfile(ready_path):
+                os.unlink(ready_path)
+            os.makedirs(tmp_dir, exist_ok=True)
+        self._world_barrier()
+
+        items: List[Tuple[str, torch.Tensor, bool]] = []
+        items.extend((n, t, False) for n, t in state_dict.items())
+        items.extend((n, t, True) for n, t in extras.items())
+
+        def _entry(item: Tuple[str, torch.Tensor, bool]) -> Tuple[str, Dict[str, Any]]:
+            name, tensor, is_extra = item
+            return name, {
+                "checksum": self._hash_tensor(tensor),
+                "size": tensor.numel() * tensor.element_size(),
+                "dtype": str(tensor.dtype),
+                "shape": list(tensor.shape),
+                "is_extra": is_extra,
+            }
+
+        manifest: Dict[str, Dict[str, Any]] = {}
+        num_workers = min(max(1, len(items)), self._hash_num_threads)
+        if num_workers <= 1:
+            for it in items:
+                name, info = _entry(it)
+                manifest[name] = info
+        else:
+            with concurrent.futures.ThreadPoolExecutor(
+                max_workers=num_workers,
+                thread_name_prefix="presharded-hash",
+            ) as ex:
+                for name, info in ex.map(_entry, items):
+                    manifest[name] = info
+
+        with open(os.path.join(tmp_dir, f"manifest_{rank:05d}.json"), "w") as f:
+            json.dump(manifest, f)
+        self._world_barrier()
+
+        if rank == 0:
+            plan = self._build_dump_plan(world_size, tmp_dir, self._max_file_bytes)
+            plan["shard_config"] = shard_config
+            with open(os.path.join(presharded_dir, self.CHECKSUM_FILENAME), "w") as f:
+                json.dump(plan, f, indent=2)
+        self._world_barrier()
+
+        with open(os.path.join(presharded_dir, self.CHECKSUM_FILENAME)) as f:
+            plan = json.load(f)
+        all_tensors = {**state_dict, **extras}
+        self._dump_files_for_rank(all_tensors, plan, rank, presharded_dir)
+        self._world_barrier()
+
+        if rank == 0:
+            shutil.rmtree(tmp_dir, ignore_errors=True)
+            ready_path = os.path.join(presharded_dir, self.READY_FILENAME)
+            with open(ready_path, "w") as f:
+                json.dump(
+                    {
+                        "plan_version": self.PLAN_VERSION,
+                        "world_size": world_size,
+                        "created_at": time.time(),
+                    },
+                    f,
+                )
+        self._world_barrier()
+
+    @staticmethod
+    def _make_filename(
+        file_id: int, rank_list: Tuple[int, ...], is_common: bool
+    ) -> str:
+        if is_common:
+            return f"model-{file_id:05d}-common.safetensor"
+        rank_str = ",".join(f"{r:03d}" for r in rank_list)
+        return f"model-{file_id:05d}-rank-{rank_str}.safetensor"
+
+    @classmethod
+    def _build_dump_plan(
+        cls, world_size: int, tmp_dir: str, max_file_bytes: int
+    ) -> Dict[str, Any]:
+        rank_to_manifest: Dict[int, Dict[str, Dict[str, Any]]] = {}
+        for r in range(world_size):
+            manifest_path = os.path.join(tmp_dir, f"manifest_{r:05d}.json")
+            try:
+                with open(manifest_path) as f:
+                    rank_to_manifest[r] = json.load(f)
+            except FileNotFoundError as e:
+                raise FileNotFoundError(
+                    f"Rank {r} did not write {manifest_path}. The presharded "
+                    "dump directory must be on a filesystem shared by all "
+                    "ranks/nodes (set presharded_path / draft_presharded_path "
+                    "to a shared path if model_path is node-local)."
+                ) from e
+
+        checksum_to_entries: Dict[str, List[Tuple[int, str, Dict[str, Any]]]] = (
+            collections.defaultdict(list)
+        )
+        name_to_is_extra: Dict[Tuple[int, str], bool] = {}
+        for r, manifest in rank_to_manifest.items():
+            for name, info in manifest.items():
+                checksum_to_entries[info["checksum"]].append((r, name, info))
+                name_to_is_extra[(r, name)] = bool(info.get("is_extra", False))
+
+        tensor_records: List[Dict[str, Any]] = []
+        for checksum, entries in checksum_to_entries.items():
+            sizes = {info["size"] for _, _, info in entries}
+            if len(sizes) != 1:
+                raise RuntimeError(
+                    f"Checksum {checksum} maps to inconsistent sizes {sizes}; "
+                    f"this indicates a hash collision or stale manifest."
+                )
+            size = next(iter(sizes))
+            ranks = sorted({r for r, _, _ in entries})
+            rank_to_names: Dict[str, List[str]] = collections.defaultdict(list)
+            for r, n, _ in entries:
+                rank_to_names[str(r)].append(n)
+            tensor_records.append(
+                {
+                    "checksum": checksum,
+                    "size": size,
+                    "rank_list": ranks,
+                    "rank_to_names": {k: sorted(v) for k, v in rank_to_names.items()},
+                }
+            )
+
+        by_rank_list: Dict[Tuple[int, ...], List[Dict[str, Any]]] = (
+            collections.defaultdict(list)
+        )
+        for rec in tensor_records:
+            by_rank_list[tuple(rec["rank_list"])].append(rec)
+
+        files: List[Dict[str, Any]] = []
+        next_file_id = 0
+        for rank_tuple, recs in by_rank_list.items():
+            recs.sort(key=lambda r: -r["size"])
+            is_common = len(rank_tuple) == world_size and rank_tuple == tuple(
+                range(world_size)
+            )
+            writer_load = {wr: 0 for wr in rank_tuple}
+            writer_records: Dict[int, List[Dict[str, Any]]] = {
+                wr: [] for wr in rank_tuple
+            }
+            for rec in recs:
+                wr = min(rank_tuple, key=lambda r: writer_load[r])
+                writer_records[wr].append(rec)
+                writer_load[wr] += rec["size"]
+
+            for wr, wr_recs in writer_records.items():
+                cur_size = 0
+                cur_tensors: List[Dict[str, Any]] = []
+                for rec in wr_recs:
+                    if cur_tensors and cur_size + rec["size"] > max_file_bytes:
+                        files.append(
+                            {
+                                "filename": cls._make_filename(
+                                    next_file_id, rank_tuple, is_common
+                                ),
+                                "writer_rank": wr,
+                                "rank_list": (None if is_common else list(rank_tuple)),
+                                "is_common": is_common,
+                                "tensors": cur_tensors,
+                            }
+                        )
+                        next_file_id += 1
+                        cur_size = 0
+                        cur_tensors = []
+                    cur_tensors.append(
+                        {
+                            "stored_key": rec["checksum"],
+                            "size": rec["size"],
+                            "rank_to_names": rec["rank_to_names"],
+                        }
+                    )
+                    cur_size += rec["size"]
+                if cur_tensors:
+                    files.append(
+                        {
+                            "filename": cls._make_filename(
+                                next_file_id, rank_tuple, is_common
+                            ),
+                            "writer_rank": wr,
+                            "rank_list": (None if is_common else list(rank_tuple)),
+                            "is_common": is_common,
+                            "tensors": cur_tensors,
+                        }
+                    )
+                    next_file_id += 1
+
+        rank_to_reads: Dict[int, List[Dict[str, Any]]] = collections.defaultdict(list)
+        for f in files:
+            for t in f["tensors"]:
+                for r_str, names in t["rank_to_names"].items():
+                    for name in names:
+                        rank_to_reads[int(r_str)].append(
+                            {
+                                "filename": f["filename"],
+                                "stored_key": t["stored_key"],
+                                "name": name,
+                                "is_extra": name_to_is_extra.get(
+                                    (int(r_str), name), False
+                                ),
+                            }
+                        )
+
+        rank_checksums: Dict[str, str] = {}
+        for r in range(world_size):
+            total = 0
+            for rec in rank_to_reads.get(r, []):
+                d = cls._fold_name_content_digest(rec["name"], rec["stored_key"])
+                total = (total + int.from_bytes(d[:8], "big")) & 0xFFFFFFFFFFFFFFFF
+            rank_checksums[str(r)] = format(total, "016x")
+
+        return {
+            "version": cls.PLAN_VERSION,
+            "world_size": world_size,
+            "files": files,
+            "rank_to_reads": {str(r): v for r, v in rank_to_reads.items()},
+            "rank_checksums": rank_checksums,
+        }
+
+    def _dump_files_for_rank(
+        self,
+        state_dict: Dict[str, torch.Tensor],
+        plan: Dict[str, Any],
+        rank: int,
+        presharded_dir: str,
+    ) -> None:
+        from safetensors.torch import save_file
+
+        for f in plan["files"]:
+            if f["writer_rank"] != rank:
+                continue
+            tensors_to_save: Dict[str, torch.Tensor] = {}
+            for t in f["tensors"]:
+                names_for_this_rank = t["rank_to_names"].get(str(rank))
+                if not names_for_this_rank:
+                    raise RuntimeError(
+                        f"writer_rank {rank} is missing tensor {t['stored_key']} "
+                        f"for file {f['filename']}; plan is inconsistent."
+                    )
+                name_for_this_rank = names_for_this_rank[0]
+                tensor = (
+                    state_dict[name_for_this_rank]
+                    .detach()
+                    .to(device="cpu", copy=False)
+                    .contiguous()
+                )
+                tensors_to_save[t["stored_key"]] = tensor
+            save_file(tensors_to_save, os.path.join(presharded_dir, f["filename"]))
+
+    @staticmethod
+    def _read_presharded_file(
+        full_path: str, stored_keys: List[str]
+    ) -> Dict[str, torch.Tensor]:
+        from safetensors.torch import safe_open
+
+        with safe_open(full_path, framework="pt") as fh:
+            return {key: fh.get_tensor(key) for key in stored_keys}
+
+    def _apply_presharded_file(
+        self,
+        *,
+        items: List[Dict[str, Any]],
+        cached: Dict[str, torch.Tensor],
+        model: nn.Module,
+        state_dict: Dict[str, torch.Tensor],
+        target_device: torch.device,
+        loaded_param_keys: set,
+        verify_hashes: List[Tuple[str, str]],
+    ) -> None:
+        if self._verify_on_load:
+            keys = list(cached.keys())
+            n_workers = min(max(1, len(keys)), self._hash_num_threads)
+
+            def _hash_one(key, _cached=cached):
+                return key, self._hash_tensor(_cached[key])
+
+            if n_workers <= 1:
+                key_to_hash = dict(_hash_one(k) for k in keys)
+            else:
+                with concurrent.futures.ThreadPoolExecutor(
+                    max_workers=n_workers,
+                    thread_name_prefix="presharded-verify",
+                ) as ex:
+                    key_to_hash = dict(ex.map(_hash_one, keys))
+            for r in items:
+                verify_hashes.append((r["name"], key_to_hash[r["stored_key"]]))
+
+        for r in items:
+            tensor = cached[r["stored_key"]]
+            if r.get("is_extra"):
+                module_path, _, attr_name = r["name"].rpartition(".")
+                module = model.get_submodule(module_path) if module_path else model
+                if hasattr(module, attr_name):
+                    try:
+                        delattr(module, attr_name)
+                    except AttributeError:
+                        pass
+                setattr(module, attr_name, tensor.to(target_device))
+                continue
+            if r["name"] not in state_dict:
+                continue
+            param_data = state_dict[r["name"]].data
+            param_shape = state_dict[r["name"]].shape
+            for dim, size in enumerate(tensor.shape):
+                if size < param_shape[dim]:
+                    param_data = param_data.narrow(dim, 0, size)
+            if tensor.shape != param_data.shape:
+                raise ValueError(
+                    f"Presharded tensor shape mismatch for '{r['name']}': "
+                    f"dumped {tuple(tensor.shape)} vs parameter slice "
+                    f"{tuple(param_data.shape)} (full param {tuple(param_shape)}). "
+                    "Re-dump with matching quant/parallel config, or set "
+                    "verify_on_load and check process_weights_after_loading."
+                )
+            param_data.copy_(tensor)
+            loaded_param_keys.add(r["name"])
+
+        cached.clear()
+        del cached
+
+    def _load_from_presharded(
+        self,
+        model_config: ModelConfig,
+        device_config: DeviceConfig,
+        presharded_dir: str,
+    ) -> nn.Module:
+        target_device = torch.device(device_config.device)
+        quant_config = _get_quantization_config(model_config, self.load_config)
+
+        with set_default_torch_dtype(model_config.dtype):
+            with target_device:
+                model = _initialize_model(model_config, self.load_config, quant_config)
+
+            for _, module in model.named_modules():
+                quant_method = getattr(module, "quant_method", None)
+                if quant_method is not None:
+                    with device_loading_context(module, target_device):
+                        quant_method.process_weights_after_loading(module)
+
+            rank, _ = self._world_rank_and_size()
+            with open(os.path.join(presharded_dir, self.CHECKSUM_FILENAME)) as f:
+                plan = json.load(f)
+            if plan.get("version") != self.PLAN_VERSION:
+                raise ValueError(
+                    f"Unsupported presharded plan version {plan.get('version')!r} "
+                    f"at {presharded_dir}; expected {self.PLAN_VERSION}."
+                )
+
+            state_dict = dict(model.state_dict())
+            reads = plan.get("rank_to_reads", {}).get(str(rank), [])
+
+            by_file: Dict[str, List[Dict[str, Any]]] = collections.defaultdict(list)
+            for r in reads:
+                by_file[r["filename"]].append(r)
+
+            loaded_param_keys: set = set()
+            verify_hashes: List[Tuple[str, str]] = []
+            for filename, items in by_file.items():
+                stored_keys = list(dict.fromkeys(r["stored_key"] for r in items))
+                cached = self._read_presharded_file(
+                    os.path.join(presharded_dir, filename), stored_keys
+                )
+                self._apply_presharded_file(
+                    items=items,
+                    cached=cached,
+                    model=model,
+                    state_dict=state_dict,
+                    target_device=target_device,
+                    loaded_param_keys=loaded_param_keys,
+                    verify_hashes=verify_hashes,
+                )
+            if torch.cuda.is_available():
+                torch.cuda.synchronize()
+
+            loaded_storages: set = set()
+            for k in loaded_param_keys:
+                t = state_dict[k]
+                if t.numel() > 0:
+                    loaded_storages.add((t.device, t.untyped_storage().data_ptr()))
+            missing = []
+            for k, t in state_dict.items():
+                if k in loaded_param_keys:
+                    continue
+                if t.numel() == 0:
+                    continue
+                storage_key = (t.device, t.untyped_storage().data_ptr())
+                if storage_key not in loaded_storages:
+                    missing.append(k)
+            if missing:
+                raise ValueError(
+                    f"Missing keys {tuple(sorted(missing))} in presharded "
+                    f"checkpoint at {presharded_dir}."
+                )
+
+            self._rebind_parameter_aliases(model)
+
+            if self._verify_on_load:
+                self._verify_rank_checksum(verify_hashes, plan, rank, presharded_dir)
+
+        self.counter_after_loading_weights = time.perf_counter()
+        return model.eval()
+
+
 class BitsAndBytesModelLoader(BaseModelLoader):
     """Model loader to load model weights with BitAndBytes quantization."""
 
@@ -3271,6 +4116,9 @@ def get_model_loader(
     if load_config.load_format == LoadFormat.SHARDED_STATE:
         return ShardedStateLoader(load_config)
 
+    if load_config.load_format == LoadFormat.PRESHARDED:
+        return PreshardedModelLoader(load_config)
+
     if load_config.load_format == LoadFormat.BITSANDBYTES:
         return BitsAndBytesModelLoader(load_config)
 

python/sglang/srt/server_args.py

@@ -113,6 +113,7 @@ LOAD_FORMAT_CHOICES = [
     "npcache",
     "dummy",
     "sharded_state",
+    "presharded",
     "gguf",
     "bitsandbytes",
     "mistral",
@@ -517,14 +518,31 @@ class ServerArgs:
             "quantization."
             '"layered" loads weights layer by layer so that one can quantize a '
             "layer before loading another to make the peak memory envelope "
-            "smaller.",
+            "smaller."
+            '"presharded" performs a normal first-time load (with quantization), '
+            "then dumps a per-rank/per-tensor sharded checkpoint with content "
+            "deduplication into "
+            "<model_path>/presharded/<parallelism+quant subfolder>/. "
+            "Subsequent runs with the same parallelism+quantization config "
+            "load directly from this presharded checkpoint and skip "
+            "re-quantization. "
+            "The dump directory must be on a shared filesystem across all "
+            "ranks/nodes. Optional model_loader_extra_config roots: "
+            "presharded_path (target) and draft_presharded_path (speculative "
+            "draft); each replaces <model_path>/presharded and still gets a "
+            "config subfolder appended. Use a writable path when model_path "
+            "is read-only (e.g. HF cache mounts).",
             choices=LOAD_FORMAT_CHOICES,
         ),
         NS("model"),
     ] = "auto"
     model_loader_extra_config: A[
         str,
-        "Extra config for model loader. This will be passed to the model loader corresponding to the chosen load_format.",
+        "Extra config for model loader. This will be passed to the model loader "
+        "corresponding to the chosen load_format. For load_format=presharded, "
+        "JSON may include presharded_path (target cache root), "
+        "draft_presharded_path (draft cache root), max_file_bytes, "
+        "hash_num_threads, and verify_on_load.",
         NS("model"),
     ] = "{}"
     trust_remote_code: A[

test/registered/unit/model_loader/test_presharded_loader.py

@@ -0,0 +1,960 @@
+"""Unit tests for PreshardedModelLoader's pure helpers.
+
+These tests exercise the deterministic pieces of the presharding algorithm
+(tensor hashing, plan construction, file naming, dedup, file-size cap, and
+per-rank workload balance) without needing a GPU or distributed setup.
+"""
+
+import json
+import os
+import tempfile
+import unittest
+from types import SimpleNamespace
+from unittest import mock
+
+import torch
+
+from sglang.srt.model_loader.loader import PreshardedModelLoader
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=10, suite="base-a-test-cpu")
+
+
+class TestPreshardedHashTensor(unittest.TestCase):
+    def test_same_content_same_hash(self):
+        a = torch.arange(100, dtype=torch.float32).reshape(10, 10)
+        b = torch.arange(100, dtype=torch.float32).reshape(10, 10)
+        self.assertEqual(
+            PreshardedModelLoader._hash_tensor(a),
+            PreshardedModelLoader._hash_tensor(b),
+        )
+
+    def test_different_content_different_hash(self):
+        a = torch.arange(100, dtype=torch.float32)
+        b = a.clone()
+        b[0] = 999
+        self.assertNotEqual(
+            PreshardedModelLoader._hash_tensor(a),
+            PreshardedModelLoader._hash_tensor(b),
+        )
+
+    def test_dtype_changes_hash(self):
+        a = torch.zeros(8, dtype=torch.float32)
+        b = torch.zeros(16, dtype=torch.float16)  # same byte content (zeros)
+        self.assertNotEqual(
+            PreshardedModelLoader._hash_tensor(a),
+            PreshardedModelLoader._hash_tensor(b),
+        )
+
+    def test_shape_changes_hash(self):
+        a = torch.zeros(16, dtype=torch.float32)
+        b = torch.zeros((4, 4), dtype=torch.float32)
+        self.assertNotEqual(
+            PreshardedModelLoader._hash_tensor(a),
+            PreshardedModelLoader._hash_tensor(b),
+        )
+
+    def test_empty_tensor_is_hashable(self):
+        a = torch.empty(0, dtype=torch.float32)
+        digest = PreshardedModelLoader._hash_tensor(a)
+        self.assertIsInstance(digest, str)
+        self.assertEqual(
+            len(digest), PreshardedModelLoader._CONTENT_HASH_HEX_LEN
+        )  # xxh3-128 hex
+
+    def test_cuda_and_cpu_digests_agree(self):
+        # Streaming D2H + host xxh3 must match a full CPU hash of the same
+        # bytes; otherwise dump/reload verify would false-fail across devices.
+        if not torch.cuda.is_available():
+            self.skipTest("CUDA required")
+        cpu = torch.arange(10_000, dtype=torch.float32)
+        gpu = cpu.cuda()
+        self.assertEqual(
+            PreshardedModelLoader._hash_tensor(cpu),
+            PreshardedModelLoader._hash_tensor(gpu),
+        )
+
+
+class TestPreshardedFilename(unittest.TestCase):
+    def test_common_filename(self):
+        self.assertEqual(
+            PreshardedModelLoader._make_filename(0, (0, 1, 2, 3), is_common=True),
+            "model-00000-common.safetensor",
+        )
+
+    def test_rank_filename_three_digit_padding(self):
+        self.assertEqual(
+            PreshardedModelLoader._make_filename(5, (1, 3, 5, 7), is_common=False),
+            "model-00005-rank-001,003,005,007.safetensor",
+        )
+
+    def test_file_id_zero_padding(self):
+        self.assertTrue(
+            PreshardedModelLoader._make_filename(42, (0,), is_common=False).startswith(
+                "model-00042-"
+            )
+        )
+
+
+class TestBuildDumpPlan(unittest.TestCase):
+    def _write_manifests(self, tmp_dir, manifests):
+        for r, m in manifests.items():
+            with open(os.path.join(tmp_dir, f"manifest_{r:05d}.json"), "w") as f:
+                json.dump(m, f)
+
+    def test_dedup_across_ranks(self):
+        # Same content (same checksum) on all 4 ranks → single file marked common.
+        with tempfile.TemporaryDirectory() as tmp:
+            shared = {
+                "checksum": "deadbeef",
+                "size": 1024,
+                "dtype": "torch.float32",
+                "shape": [256],
+            }
+            manifests = {r: {"shared.weight": shared} for r in range(4)}
+            self._write_manifests(tmp, manifests)
+            plan = PreshardedModelLoader._build_dump_plan(
+                world_size=4, tmp_dir=tmp, max_file_bytes=10**12
+            )
+        self.assertEqual(len(plan["files"]), 1)
+        self.assertTrue(plan["files"][0]["is_common"])
+        self.assertIn("common", plan["files"][0]["filename"])
+        # Each rank should still have a read entry pointing at the file.
+        for r in range(4):
+            reads = plan["rank_to_reads"][str(r)]
+            self.assertEqual(len(reads), 1)
+            self.assertEqual(reads[0]["name"], "shared.weight")
+            self.assertEqual(reads[0]["filename"], plan["files"][0]["filename"])
+            self.assertEqual(reads[0]["stored_key"], "deadbeef")
+        self.assertIn("rank_checksums", plan)
+        self.assertEqual(set(plan["rank_checksums"].keys()), {"0", "1", "2", "3"})
+
+    def test_per_rank_unique_tensors(self):
+        # Each rank has its own tensor (different content). 4 distinct files,
+        # filenames should be rank-{rrr}.
+        with tempfile.TemporaryDirectory() as tmp:
+            manifests = {
+                r: {
+                    "layer.weight": {
+                        "checksum": f"hash_{r}",
+                        "size": 2048,
+                        "dtype": "torch.float32",
+                        "shape": [512],
+                    }
+                }
+                for r in range(4)
+            }
+            self._write_manifests(tmp, manifests)
+            plan = PreshardedModelLoader._build_dump_plan(
+                world_size=4, tmp_dir=tmp, max_file_bytes=10**12
+            )
+        self.assertEqual(len(plan["files"]), 4)
+        for f in plan["files"]:
+            self.assertFalse(f["is_common"])
+            self.assertEqual(len(f["rank_list"]), 1)
+            # Writer is the only rank in the rank_list.
+            self.assertEqual(f["writer_rank"], f["rank_list"][0])
+            self.assertIn(f"-rank-{f['rank_list'][0]:03d}.safetensor", f["filename"])
+
+    def test_partial_share_has_correct_rank_list(self):
+        # Tensor shared by ranks 1,3,5,7 only.
+        with tempfile.TemporaryDirectory() as tmp:
+            shared = {
+                "checksum": "shared_hash",
+                "size": 1024,
+                "dtype": "torch.float32",
+                "shape": [256],
+            }
+            manifests = {
+                r: ({"x": shared} if r in (1, 3, 5, 7) else {}) for r in range(8)
+            }
+            self._write_manifests(tmp, manifests)
+            plan = PreshardedModelLoader._build_dump_plan(
+                world_size=8, tmp_dir=tmp, max_file_bytes=10**12
+            )
+        self.assertEqual(len(plan["files"]), 1)
+        f = plan["files"][0]
+        self.assertFalse(f["is_common"])
+        self.assertEqual(f["rank_list"], [1, 3, 5, 7])
+        self.assertIn(f["writer_rank"], (1, 3, 5, 7))
+        self.assertEqual(f["filename"], "model-00000-rank-001,003,005,007.safetensor")
+        # Only the 4 sharing ranks have read entries.
+        for r in range(8):
+            reads = plan["rank_to_reads"].get(str(r), [])
+            if r in (1, 3, 5, 7):
+                self.assertEqual(len(reads), 1)
+            else:
+                self.assertEqual(len(reads), 0)
+
+    def test_max_file_size_caps_files(self):
+        # Two tensors of 1 MiB each shared by all 2 ranks; cap = 1.5 MiB →
+        # two files.
+        with tempfile.TemporaryDirectory() as tmp:
+            t1 = {
+                "checksum": "h1",
+                "size": 1024 * 1024,
+                "dtype": "torch.float32",
+                "shape": [256, 1024],
+            }
+            t2 = dict(t1, checksum="h2")
+            manifests = {0: {"a": t1, "b": t2}, 1: {"a": t1, "b": t2}}
+            self._write_manifests(tmp, manifests)
+            plan = PreshardedModelLoader._build_dump_plan(
+                world_size=2,
+                tmp_dir=tmp,
+                max_file_bytes=int(1.5 * 1024 * 1024),
+            )
+        # Both tensors share the same rank_list (0,1). With balanced packing,
+        # each writer (rank 0 and rank 1) gets one tensor → 2 files.
+        self.assertEqual(len(plan["files"]), 2)
+        for f in plan["files"]:
+            self.assertEqual(len(f["tensors"]), 1)
+            # rank_list is full world ⇒ marked common
+            self.assertTrue(f["is_common"])
+
+    def test_workload_balanced_within_rank_list(self):
+        # 4 same-size tensors all shared by ranks (0,1,2,3) → with balanced
+        # packing each writer rank should get exactly one tensor.
+        with tempfile.TemporaryDirectory() as tmp:
+            tensors = {
+                f"t{i}": {
+                    "checksum": f"h{i}",
+                    "size": 1024,
+                    "dtype": "torch.float32",
+                    "shape": [256],
+                }
+                for i in range(4)
+            }
+            manifests = {r: dict(tensors) for r in range(4)}
+            self._write_manifests(tmp, manifests)
+            plan = PreshardedModelLoader._build_dump_plan(
+                world_size=4,
+                tmp_dir=tmp,
+                max_file_bytes=10**12,
+            )
+        # 4 tensors, 4 writers, 4 files (one per writer).
+        self.assertEqual(len(plan["files"]), 4)
+        writers = sorted(f["writer_rank"] for f in plan["files"])
+        self.assertEqual(writers, [0, 1, 2, 3])
+
+    def test_round_trip_dump_and_read_back(self):
+        # End-to-end on disk for world_size=1: build manifests for tensors,
+        # construct plan, write safetensors per the plan, read back and
+        # verify both checksums and bit-identity.
+        from safetensors.torch import safe_open, save_file
+
+        with tempfile.TemporaryDirectory() as tmp:
+            tensors = {
+                "embed.weight": torch.arange(64, dtype=torch.float32).reshape(8, 8),
+                "norm.weight": torch.full((16,), 0.5, dtype=torch.float32),
+                "head.weight": torch.linspace(-1, 1, 32, dtype=torch.float32),
+            }
+            manifest_dir = os.path.join(tmp, "manifests")
+            presharded_dir = os.path.join(tmp, "presharded")
+            os.makedirs(manifest_dir)
+            os.makedirs(presharded_dir)
+
+            checksums = {
+                name: PreshardedModelLoader._hash_tensor(t)
+                for name, t in tensors.items()
+            }
+            manifest = {
+                name: {
+                    "checksum": checksums[name],
+                    "size": t.numel() * t.element_size(),
+                    "dtype": str(t.dtype),
+                    "shape": list(t.shape),
+                }
+                for name, t in tensors.items()
+            }
+            with open(os.path.join(manifest_dir, "manifest_00000.json"), "w") as f:
+                json.dump(manifest, f)
+
+            plan = PreshardedModelLoader._build_dump_plan(
+                world_size=1, tmp_dir=manifest_dir, max_file_bytes=10**12
+            )
+
+            # Single rank → all tensors live in one common file.
+            self.assertEqual(len(plan["files"]), 1)
+            f = plan["files"][0]
+            self.assertTrue(f["is_common"])
+
+            # Write the file with stored_keys mapped to tensor content.
+            to_save = {}
+            for t_entry in f["tensors"]:
+                name = t_entry["rank_to_names"]["0"][0]
+                to_save[t_entry["stored_key"]] = tensors[name]
+            save_file(to_save, os.path.join(presharded_dir, f["filename"]))
+
+            # Read back: verify each tensor's checksum and content.
+            with safe_open(
+                os.path.join(presharded_dir, f["filename"]), framework="pt"
+            ) as fh:
+                for r in plan["rank_to_reads"]["0"]:
+                    loaded = fh.get_tensor(r["stored_key"])
+                    self.assertEqual(
+                        PreshardedModelLoader._hash_tensor(loaded), r["stored_key"]
+                    )
+                    torch.testing.assert_close(loaded, tensors[r["name"]])
+
+    def test_dedup_with_multiple_names_per_rank(self):
+        # Same checksum can appear under MULTIPLE param names on the same
+        # rank (e.g., k_scale and v_scale both default to 1.0). Both names
+        # must end up in rank_to_reads.
+        with tempfile.TemporaryDirectory() as tmp:
+            shared = {
+                "checksum": "scale_hash",
+                "size": 4,
+                "dtype": "torch.float32",
+                "shape": [],
+            }
+            manifests = {
+                0: {
+                    "layers.0.attn.k_scale": shared,
+                    "layers.0.attn.v_scale": shared,
+                    "layers.1.attn.k_scale": shared,
+                },
+                1: {
+                    "layers.0.attn.k_scale": shared,
+                    "layers.0.attn.v_scale": shared,
+                    "layers.1.attn.k_scale": shared,
+                },
+            }
+            self._write_manifests(tmp, manifests)
+            plan = PreshardedModelLoader._build_dump_plan(
+                world_size=2, tmp_dir=tmp, max_file_bytes=10**12
+            )
+        # All 6 (rank,name) pairs must be readable, even though there is one
+        # underlying tensor stored on disk.
+        self.assertEqual(len(plan["files"]), 1)
+        for r in (0, 1):
+            reads = plan["rank_to_reads"][str(r)]
+            names = sorted(rd["name"] for rd in reads)
+            self.assertEqual(
+                names,
+                [
+                    "layers.0.attn.k_scale",
+                    "layers.0.attn.v_scale",
+                    "layers.1.attn.k_scale",
+                ],
+            )
+            # All point at the same stored_key (deduplicated content).
+            self.assertEqual({rd["stored_key"] for rd in reads}, {"scale_hash"})
+
+    def test_collision_size_mismatch_raises(self):
+        # Same checksum but different sizes ⇒ plan builder rejects.
+        with tempfile.TemporaryDirectory() as tmp:
+            manifests = {
+                0: {
+                    "x": {
+                        "checksum": "same",
+                        "size": 1024,
+                        "dtype": "torch.float32",
+                        "shape": [256],
+                    }
+                },
+                1: {
+                    "x": {
+                        "checksum": "same",
+                        "size": 2048,
+                        "dtype": "torch.float32",
+                        "shape": [512],
+                    }
+                },
+            }
+            self._write_manifests(tmp, manifests)
+            with self.assertRaises(RuntimeError):
+                PreshardedModelLoader._build_dump_plan(
+                    world_size=2, tmp_dir=tmp, max_file_bytes=10**12
+                )
+
+    def test_rank_checksum_deterministic(self):
+        # rank_checksums must be reproducible from the same manifest input
+        # and depend on (name, content-SHA) pairs of every tensor a rank
+        # owns. Permuting the manifest's insertion order must not change
+        # the rank checksum.
+        with tempfile.TemporaryDirectory() as tmp_a, tempfile.TemporaryDirectory() as tmp_b:
+            base_entries = {
+                "alpha.weight": {
+                    "checksum": "h_alpha",
+                    "size": 16,
+                    "dtype": "torch.float32",
+                    "shape": [4],
+                },
+                "beta.weight": {
+                    "checksum": "h_beta",
+                    "size": 16,
+                    "dtype": "torch.float32",
+                    "shape": [4],
+                },
+            }
+            self._write_manifests(tmp_a, {0: dict(base_entries)})
+            # Insertion-order-permuted copy.
+            permuted = {k: base_entries[k] for k in reversed(list(base_entries))}
+            self._write_manifests(tmp_b, {0: permuted})
+            plan_a = PreshardedModelLoader._build_dump_plan(
+                world_size=1, tmp_dir=tmp_a, max_file_bytes=10**12
+            )
+            plan_b = PreshardedModelLoader._build_dump_plan(
+                world_size=1, tmp_dir=tmp_b, max_file_bytes=10**12
+            )
+            self.assertEqual(plan_a["rank_checksums"], plan_b["rank_checksums"])
+
+    def test_rank_checksum_distinguishes_content(self):
+        # Changing one tensor's content-SHA must change the rank checksum.
+        with tempfile.TemporaryDirectory() as tmp_a, tempfile.TemporaryDirectory() as tmp_b:
+            entries_a = {
+                "x.weight": {
+                    "checksum": "ha",
+                    "size": 16,
+                    "dtype": "torch.float32",
+                    "shape": [4],
+                },
+            }
+            entries_b = {
+                "x.weight": {
+                    "checksum": "hb",  # different content
+                    "size": 16,
+                    "dtype": "torch.float32",
+                    "shape": [4],
+                },
+            }
+            self._write_manifests(tmp_a, {0: entries_a})
+            self._write_manifests(tmp_b, {0: entries_b})
+            plan_a = PreshardedModelLoader._build_dump_plan(
+                world_size=1, tmp_dir=tmp_a, max_file_bytes=10**12
+            )
+            plan_b = PreshardedModelLoader._build_dump_plan(
+                world_size=1, tmp_dir=tmp_b, max_file_bytes=10**12
+            )
+            self.assertNotEqual(
+                plan_a["rank_checksums"]["0"],
+                plan_b["rank_checksums"]["0"],
+            )
+
+    def test_presharded_ready_sentinel(self):
+        # Loader treats a dir as a valid presharded ckpt only when the
+        # READY sentinel exists. A bare checksum.json (or partial files)
+        # is NOT enough.
+        with tempfile.TemporaryDirectory() as tmp:
+            self.assertFalse(PreshardedModelLoader._presharded_ready(tmp))
+            # checksum.json alone is not sufficient.
+            with open(
+                os.path.join(tmp, PreshardedModelLoader.CHECKSUM_FILENAME), "w"
+            ) as f:
+                f.write("{}")
+            self.assertFalse(PreshardedModelLoader._presharded_ready(tmp))
+            # READY makes it ready.
+            with open(
+                os.path.join(tmp, PreshardedModelLoader.READY_FILENAME), "w"
+            ) as f:
+                f.write("{}")
+            self.assertTrue(PreshardedModelLoader._presharded_ready(tmp))
+
+    def test_separate_presharded_path_overrides_for_target_and_draft(self):
+        # Target and draft get distinct roots via presharded_path vs
+        # draft_presharded_path. Failure mode if this regresses: draft
+        # re-dump collides with / wipes target READY under a shared path.
+        loader = object.__new__(PreshardedModelLoader)
+        with tempfile.TemporaryDirectory() as tmp:
+            target_root = os.path.join(tmp, "target_cache")
+            draft_root = os.path.join(tmp, "draft_cache")
+            loader._presharded_path_override = target_root
+            loader._draft_presharded_path_override = draft_root
+            cfg = {"tp": 8, "structural_signature": "same_sig"}
+            sub = loader._build_subfolder_name(cfg)
+            target_mc = SimpleNamespace(model_path="/models/dsv3", is_draft_model=False)
+            draft_mc = SimpleNamespace(model_path="/models/dsv3", is_draft_model=True)
+            target_dir = loader._presharded_dir(target_mc, cfg)
+            draft_dir = loader._presharded_dir(draft_mc, cfg)
+            self.assertEqual(target_dir, os.path.join(target_root, sub))
+            self.assertEqual(draft_dir, os.path.join(draft_root, sub))
+            self.assertNotEqual(target_dir, draft_dir)
+
+            # Target-only override: draft must not fall back into the target
+            # root (same model_path is common for DeepSeek MTP).
+            loader._draft_presharded_path_override = None
+            draft_fallback = loader._presharded_dir(draft_mc, cfg)
+            self.assertEqual(
+                draft_fallback,
+                os.path.join(
+                    "/models/dsv3",
+                    PreshardedModelLoader.DEFAULT_SUBDIR,
+                    sub,
+                ),
+            )
+            self.assertFalse(draft_fallback.startswith(target_root))
+
+            # No overrides: both use model_path/presharded/<subfolder>.
+            loader._presharded_path_override = None
+            self.assertEqual(
+                loader._presharded_dir(target_mc, cfg),
+                os.path.join(
+                    "/models/dsv3",
+                    PreshardedModelLoader.DEFAULT_SUBDIR,
+                    sub,
+                ),
+            )
+
+    def test_apply_shape_mismatch_raises(self):
+        # Reload must not silently copy_ into a wrong layout when process
+        # shapes and dumped tensors disagree (previously only warned).
+        # Use a *larger* dumped tensor so prefix-narrow does not paper over it.
+        loader = object.__new__(PreshardedModelLoader)
+        loader._verify_on_load = False
+        param = torch.nn.Parameter(torch.zeros(2, 2))
+        state_dict = {"w": param}
+        items = [{"name": "w", "stored_key": "w", "is_extra": False}]
+        cached = {"w": torch.ones(4, 4)}
+        loaded: set = set()
+        with self.assertRaises(ValueError) as ctx:
+            loader._apply_presharded_file(
+                items=items,
+                cached=cached,
+                model=torch.nn.Module(),
+                state_dict=state_dict,
+                target_device=torch.device("cpu"),
+                loaded_param_keys=loaded,
+                verify_hashes=[],
+            )
+        self.assertIn("shape mismatch", str(ctx.exception).lower())
+
+    def test_build_dump_plan_missing_manifest_mentions_shared_fs(self):
+        # Multi-node without a shared dump dir fails at plan build with a
+        # clear message (not a bare FileNotFoundError path).
+        with tempfile.TemporaryDirectory() as tmp:
+            with self.assertRaises(FileNotFoundError) as ctx:
+                PreshardedModelLoader._build_dump_plan(
+                    world_size=2, tmp_dir=tmp, max_file_bytes=1024
+                )
+            msg = str(ctx.exception)
+            self.assertIn("Rank 0", msg)
+            self.assertIn("shared", msg.lower())
+
+    def test_ensure_presharded_dir_writable_rejects_readonly(self):
+        # Guard against spending a full source load before discovering a
+        # read-only dump root (HF cache mounts). Mock OSError because root
+        # can often still write to mode-0555 dirs under CAP_DAC_OVERRIDE.
+        loader = object.__new__(PreshardedModelLoader)
+        with mock.patch.object(
+            loader, "_world_rank_and_size", return_value=(0, 1)
+        ), mock.patch.object(loader, "_world_barrier"), mock.patch(
+            "os.makedirs", side_effect=OSError("Read-only file system")
+        ):
+            with self.assertRaises(RuntimeError) as ctx:
+                loader._ensure_presharded_dir_writable("/ro/presharded")
+            self.assertIn("not writable", str(ctx.exception).lower())
+            self.assertIn("presharded_path", str(ctx.exception))
+
+    def test_ensure_presharded_dir_writable_ok_rank0(self):
+        loader = object.__new__(PreshardedModelLoader)
+        with tempfile.TemporaryDirectory() as tmp:
+            leaf = os.path.join(tmp, "TP-8-sig-test")
+            with mock.patch.object(
+                loader, "_world_rank_and_size", return_value=(0, 1)
+            ), mock.patch.object(loader, "_world_barrier") as barrier:
+                loader._ensure_presharded_dir_writable(leaf)
+            self.assertTrue(os.path.isdir(leaf))
+            barrier.assert_called_once()
+
+
+class TestStructuralSignature(unittest.TestCase):
+    """The structural signature is the long-term fix for the underlying
+    problem `moe_dense_tp_size` was an instance of: instead of hand-
+    enumerating every parallelism knob that might change a rank's tensor
+    shapes, hash the (name, shape, dtype) of every parameter in a
+    meta-device model skeleton built under the live parallel state. Any
+    future sharding knob that changes shapes is automatically caught
+    without touching `_build_subfolder_name`."""
+
+    def test_hash_changes_with_shape_dtype_or_added_tensor(self):
+        # Regression: a future sharding knob that changes shapes/dtypes or
+        # adds per-rank buffers must change the digest, otherwise cache
+        # collisions load wrong weights. One case covers the three axes that
+        # `_hash_structural_signature` is responsible for.
+        base = [("a.weight", (4, 4), "torch.float32")]
+        shape = [("a.weight", (4, 8), "torch.float32")]
+        dtype = [("a.weight", (4, 4), "torch.float16")]
+        extra = [
+            ("a.weight", (4, 4), "torch.float32"),
+            ("a.extra_buf", (4,), "torch.float32"),
+        ]
+        h = PreshardedModelLoader._hash_structural_signature
+        self.assertNotEqual(h(base), h(shape))
+        self.assertNotEqual(h(base), h(dtype))
+        self.assertNotEqual(h(base), h(extra))
+
+    def test_local_signature_sorts_state_dict_order(self):
+        # Production path sorts state_dict items before hashing. If that
+        # sorted(...) is dropped, two identical models whose state_dict
+        # iteration order differs would get different signatures and thrash
+        # the cache. Guard the sort, not just the pure hash helper.
+        import torch.nn as nn
+
+        class OrderedModule(nn.Module):
+            def __init__(self, order):
+                super().__init__()
+                for name in order:
+                    self.register_parameter(
+                        name, nn.Parameter(torch.zeros(2, 2), requires_grad=False)
+                    )
+
+        def _init_stub(model_config, load_config, quant_config):
+            return OrderedModule(model_config.order)
+
+        loader = object.__new__(PreshardedModelLoader)
+        loader.load_config = SimpleNamespace()
+        with mock.patch(
+            "sglang.srt.model_loader.loader._get_quantization_config",
+            return_value=None,
+        ), mock.patch(
+            "sglang.srt.model_loader.loader._initialize_model",
+            side_effect=_init_stub,
+        ), mock.patch(
+            "sglang.srt.model_loader.loader.set_default_torch_dtype",
+            return_value=mock.MagicMock(__enter__=mock.Mock(), __exit__=mock.Mock()),
+        ):
+            sig_ab = loader._compute_local_structural_signature(
+                SimpleNamespace(
+                    quantization=None, dtype=torch.float32, order=["a", "b"]
+                )
+            )
+            sig_ba = loader._compute_local_structural_signature(
+                SimpleNamespace(
+                    quantization=None, dtype=torch.float32, order=["b", "a"]
+                )
+            )
+        self.assertIsNotNone(sig_ab)
+        self.assertEqual(sig_ab, sig_ba)
+
+    def test_rank_invariant_signature_aggregates_per_rank_locals(self):
+        # Under PP, ranks build different local digests. The shared cache key
+        # must still agree across ranks: all-gather then hash the ordered
+        # list. Without this, each PP stage would pick a different
+        # presharded subfolder and the multi-rank dump protocol breaks.
+        fake_group = mock.Mock()
+        fake_group.world_size = 2
+        # Simulate two ranks each calling with their own local sig; both
+        # must see the same gathered list and thus the same aggregate.
+        fake_group.all_gather_object.side_effect = lambda local: ["sig-pp0", "sig-pp1"]
+
+        with mock.patch(
+            "sglang.srt.distributed.get_world_group", return_value=fake_group
+        ):
+            agg_from_rank0 = (
+                PreshardedModelLoader._make_rank_invariant_structural_signature(
+                    "sig-pp0"
+                )
+            )
+            agg_from_rank1 = (
+                PreshardedModelLoader._make_rank_invariant_structural_signature(
+                    "sig-pp1"
+                )
+            )
+        self.assertIsNotNone(agg_from_rank0)
+        self.assertEqual(agg_from_rank0, agg_from_rank1)
+        # Changing any rank's local contribution must change the aggregate.
+        fake_group.all_gather_object.side_effect = lambda local: [
+            "sig-pp0",
+            "sig-pp1-changed",
+        ]
+        with mock.patch(
+            "sglang.srt.distributed.get_world_group", return_value=fake_group
+        ):
+            agg_changed = (
+                PreshardedModelLoader._make_rank_invariant_structural_signature(
+                    "sig-pp0"
+                )
+            )
+        self.assertNotEqual(agg_from_rank0, agg_changed)
+
+    def test_compute_structural_signature_returns_none_on_failure(self):
+        # No self.load_config (and no real model class behind
+        # SimpleNamespace) means _get_quantization_config / _initialize_model
+        # will raise; this must be swallowed and return None, never raise,
+        # since a model class that can't be built on meta device must not
+        # break the overall load.
+        loader = object.__new__(PreshardedModelLoader)
+        model_config = SimpleNamespace(quantization=None)
+        self.assertIsNone(loader._compute_structural_signature(model_config))
+
+    def test_compute_structural_signature_picks_up_meta_model_shapes(self):
+        # End-to-end on a tiny real nn.Module standing in for the model
+        # class, to prove the meta-device construction + hashing wiring
+        # actually reflects shapes that depend on the live parallel state
+        # (here simulated via a width captured at construction time).
+        import torch.nn as nn
+
+        class FakeModelLoader(PreshardedModelLoader):
+            def __init__(self, width):
+                self._width = width
+                self.load_config = SimpleNamespace()
+
+        def _initialize_model_stub(model_config, load_config, quant_config):
+            return nn.Linear(4, model_config.width, bias=False)
+
+        with mock.patch(
+            "sglang.srt.model_loader.loader._get_quantization_config",
+            return_value=None,
+        ), mock.patch(
+            "sglang.srt.model_loader.loader._initialize_model",
+            side_effect=_initialize_model_stub,
+        ), mock.patch(
+            "sglang.srt.model_loader.loader.set_default_torch_dtype",
+            return_value=mock.MagicMock(__enter__=mock.Mock(), __exit__=mock.Mock()),
+        ):
+            loader_narrow = FakeModelLoader(width=2)
+            loader_wide = FakeModelLoader(width=8)
+            model_config_narrow = SimpleNamespace(
+                quantization=None, dtype=torch.float32, width=2
+            )
+            model_config_wide = SimpleNamespace(
+                quantization=None, dtype=torch.float32, width=8
+            )
+            sig_narrow = loader_narrow._compute_structural_signature(
+                model_config_narrow
+            )
+            sig_wide = loader_wide._compute_structural_signature(model_config_wide)
+        self.assertIsNotNone(sig_narrow)
+        self.assertIsNotNone(sig_wide)
+        self.assertNotEqual(sig_narrow, sig_wide)
+
+    def test_meta_rope_cache_cleared_even_on_failure(self):
+        # If _initialize_model partially populates _ROPE_DICT with meta-device
+        # entries before raising, _compute_structural_signature must still
+        # clean them up (via finally), otherwise the real model init reuses
+        # the meta module and fails with "Cannot copy out of meta tensor".
+        import torch.nn as nn
+
+        from sglang.srt.layers.rotary_embedding.factory import _ROPE_DICT
+
+        # Plant a fake meta-device rotary module in the global cache.
+        fake_key = ("_test_meta_rope_cleanup_sentinel",)
+        with torch.device("meta"):
+            fake_module = nn.Linear(4, 4)
+        _ROPE_DICT[fake_key] = fake_module
+
+        loader = object.__new__(PreshardedModelLoader)
+        loader.load_config = SimpleNamespace()
+
+        with mock.patch(
+            "sglang.srt.model_loader.loader._get_quantization_config",
+            return_value=None,
+        ), mock.patch(
+            "sglang.srt.model_loader.loader._initialize_model",
+            side_effect=RuntimeError("simulated init failure"),
+        ), mock.patch(
+            "sglang.srt.model_loader.loader.set_default_torch_dtype",
+            return_value=mock.MagicMock(__enter__=mock.Mock(), __exit__=mock.Mock()),
+        ):
+            result = loader._compute_structural_signature(
+                SimpleNamespace(quantization=None, dtype=torch.float32)
+            )
+
+        self.assertIsNone(result)
+        self.assertNotIn(fake_key, _ROPE_DICT)
+
+
+class TestShardConfig(unittest.TestCase):
+    """Bookkeeping guards for the enumerated cache-key fields and the
+    load-time match path. Dropping a field from `_collect_shard_config` or
+    breaking `_shard_config_matches` would silently collide caches; these
+    cases pin the failure modes the PR is meant to prevent."""
+
+    def _base_config(self, **overrides):
+        cfg = {
+            "tp": 8,
+            "dp": 1,
+            "ep": 1,
+            "pp": 1,
+            "moe_dense_tp_size": None,
+            "moe_dp_size": 1,
+            "enable_dp_lm_head": False,
+            "enable_fp32_lm_head": False,
+            "quantization": None,
+            "model_dtype": "torch.bfloat16",
+            "ep_num_redundant_experts": 0,
+            "enable_eplb": False,
+            "init_expert_location": "trivial",
+            "structural_signature": "deadbeef",
+        }
+        cfg.update(overrides)
+        return cfg
+
+    def test_collect_shard_config_includes_required_keys(self):
+        # Registry completeness: dropping a key from the dict literal in
+        # `_collect_shard_config` is the exact failure mode that left
+        # moe_dense_tp_size / LM-head flags out of the cache key before.
+        loader = object.__new__(PreshardedModelLoader)
+        server_args = SimpleNamespace(
+            moe_dense_tp_size=1,
+            moe_dp_size=2,
+            enable_dp_lm_head=True,
+            enable_fp32_lm_head=True,
+            ep_num_redundant_experts=4,
+            enable_eplb=True,
+            init_expert_location="trivial",
+        )
+        model_config = SimpleNamespace(quantization="fp8", dtype=torch.bfloat16)
+        required = {
+            "tp",
+            "dp",
+            "ep",
+            "pp",
+            "moe_dense_tp_size",
+            "moe_dp_size",
+            "enable_dp_lm_head",
+            "enable_fp32_lm_head",
+            "quantization",
+            "model_dtype",
+            "ep_num_redundant_experts",
+            "enable_eplb",
+            "init_expert_location",
+            "structural_signature",
+        }
+        parallel = SimpleNamespace(tp_size=8, moe_dp_size=2, moe_ep_size=4, pp_size=1)
+        with mock.patch(
+            "sglang.srt.model_loader.loader.get_server_args",
+            return_value=server_args,
+        ), mock.patch(
+            "sglang.srt.model_loader.loader.get_parallel",
+            return_value=parallel,
+        ), mock.patch.object(
+            loader, "_compute_structural_signature", return_value="sig16"
+        ):
+            cfg = loader._collect_shard_config(model_config)
+        self.assertEqual(required, set(cfg.keys()))
+        self.assertEqual(cfg["tp"], 8)
+        self.assertEqual(cfg["dp"], 2)
+        self.assertEqual(cfg["ep"], 4)
+        self.assertEqual(cfg["pp"], 1)
+        self.assertEqual(cfg["moe_dense_tp_size"], 1)
+        self.assertEqual(cfg["moe_dp_size"], 2)
+        self.assertTrue(cfg["enable_dp_lm_head"])
+        self.assertTrue(cfg["enable_fp32_lm_head"])
+        self.assertEqual(cfg["init_expert_location"], "trivial")
+
+    def test_enumerated_fields_change_subfolder_hash(self):
+        # Each enumerated content/shape knob must feed the subfolder name.
+        # A silent drop from `_collect_shard_config` would leave this red.
+        loader = object.__new__(PreshardedModelLoader)
+        base = self._base_config()
+        base_name = loader._build_subfolder_name(base)
+        field_variants = {
+            "moe_dense_tp_size": 1,
+            "moe_dp_size": 2,
+            "enable_dp_lm_head": True,
+            "enable_fp32_lm_head": True,
+            "ep_num_redundant_experts": 8,
+            "enable_eplb": True,
+            "init_expert_location": "file:map.json:sha1:abcd",
+            "structural_signature": "cafebabe",
+            "quantization": "fp8",
+        }
+        for field, value in field_variants.items():
+            with self.subTest(field=field):
+                other = self._base_config(**{field: value})
+                other_name = loader._build_subfolder_name(other)
+                self.assertNotEqual(
+                    base_name,
+                    other_name,
+                    f"changing {field} must change the cache subfolder name",
+                )
+
+    def test_shard_config_matches_equality_and_missing(self):
+        # Match must require exact stored equality; missing/mismatched
+        # shard_config is a cache miss (never raises) so reload re-dumps.
+        loader = object.__new__(PreshardedModelLoader)
+        cfg = self._base_config()
+        with tempfile.TemporaryDirectory() as tmp:
+            # No checksum.json → miss.
+            self.assertFalse(loader._shard_config_matches(tmp, cfg))
+            # Matching stored config → hit.
+            with open(
+                os.path.join(tmp, PreshardedModelLoader.CHECKSUM_FILENAME), "w"
+            ) as f:
+                json.dump({"shard_config": cfg}, f)
+            self.assertTrue(loader._shard_config_matches(tmp, cfg))
+            # One field differs → miss.
+            mismatched = self._base_config(moe_dense_tp_size=1)
+            self.assertFalse(loader._shard_config_matches(tmp, mismatched))
+            # Plan without shard_config key → miss (upgrade path).
+            with open(
+                os.path.join(tmp, PreshardedModelLoader.CHECKSUM_FILENAME), "w"
+            ) as f:
+                json.dump({"version": 1}, f)
+            self.assertFalse(loader._shard_config_matches(tmp, cfg))
+
+    def test_init_expert_location_hashes_file_contents(self):
+        # Overwriting the same path with a different expert map must bust
+        # the cache; keying on the path alone would silently reuse weights.
+        with tempfile.TemporaryDirectory() as tmp:
+            path = os.path.join(tmp, "experts.json")
+            with open(path, "w") as f:
+                json.dump({"logical_count": [[1, 0], [0, 1]]}, f)
+            key_a = PreshardedModelLoader._normalize_init_expert_location(path)
+            with open(path, "w") as f:
+                json.dump({"logical_count": [[0, 1], [1, 0]]}, f)
+            key_b = PreshardedModelLoader._normalize_init_expert_location(path)
+        self.assertNotEqual(key_a, key_b)
+        self.assertTrue(key_a.startswith("file:experts.json:sha1:"))
+        self.assertEqual(
+            PreshardedModelLoader._normalize_init_expert_location("trivial"),
+            "trivial",
+        )
+
+    def test_redump_clears_ready_before_rewrite(self):
+        # Config-mismatch re-dump into an already-ready dir must drop READY
+        # before mutating files, otherwise concurrent readers can observe
+        # READY + partial checksum/safetensors.
+        loader = object.__new__(PreshardedModelLoader)
+        loader._hash_num_threads = 1
+        loader._max_file_bytes = 10**12
+
+        with tempfile.TemporaryDirectory() as tmp:
+            ready_path = os.path.join(tmp, PreshardedModelLoader.READY_FILENAME)
+            with open(ready_path, "w") as f:
+                f.write("{}")
+            self.assertTrue(os.path.isfile(ready_path))
+
+            # Force rank 0 / world 1 so the method runs the rank-0 prologue
+            # and then fails early on empty state (no need for full dump).
+            with mock.patch.object(
+                PreshardedModelLoader,
+                "_world_rank_and_size",
+                return_value=(0, 1),
+            ), mock.patch.object(
+                PreshardedModelLoader, "_world_barrier", return_value=None
+            ), mock.patch.object(
+                PreshardedModelLoader,
+                "_build_dump_plan",
+                return_value={
+                    "version": PreshardedModelLoader.PLAN_VERSION,
+                    "files": [],
+                    "rank_to_reads": {"0": []},
+                    "rank_checksums": {"0": "0"},
+                },
+            ), mock.patch.object(
+                PreshardedModelLoader, "_dump_files_for_rank", return_value=None
+            ):
+                loader._dump_state_to_disk(
+                    state_dict={},
+                    extras={},
+                    presharded_dir=tmp,
+                    shard_config=self._base_config(),
+                )
+            # Dump rewrote READY at the end; the critical property is that
+            # the prologue unlinked the *previous* READY before rewriting
+            # checksum.json. Assert checksum was written and READY exists
+            # only as the fresh sentinel from this dump.
+            self.assertTrue(os.path.isfile(ready_path))
+            with open(ready_path) as f:
+                sentinel = json.load(f)
+            self.assertIn("plan_version", sentinel)
+            with open(os.path.join(tmp, PreshardedModelLoader.CHECKSUM_FILENAME)) as f:
+                plan = json.load(f)
+            self.assertEqual(plan["shard_config"]["tp"], 8)
+
+
+if __name__ == "__main__":
+    unittest.main()