Weight loading, sharded loaders, adding a model

Cold start is the one part of an inference engine where the GPU is idle by construction. The model runner is built, the KV cache is sized, the CUDA graphs are ready to capture — and none of it matters until bytes from a checkpoint are sitting in the right memory of the right device. Llama-3-8B is 16.06 GB of bf16 weights; Llama-3-70B is 131 GiB (141.1 GB). Those bytes must cross whatever medium the checkpoint lives on, and storage is the slowest link in the serving stack — three orders of magnitude below the 3.35 TB/s of HBM those weights sit in once they arrive. Note that this chapter's figure is the resident checkpoint, every parameter of it, because every parameter has to be read off disk and placed. It is not the figure that sets the decode floor: a decode step streams only 7.50 B of Llama-3-8B's 8.03 B parameters, because the embedding table is gathered rather than multiplied, which is why §0.4 divides 15.01 GB rather than 16.06 GB by the bandwidth and gets 4.48 ms. Do not carry one number into the other's derivation.

Work the floor from link rates, which are pure arithmetic. A PCIe Gen4 x4 NVMe drive signals at 16 GT/s per lane over 4 lanes with 128b/130b encoding: $16 \times 4 \times \frac{128}{130} / 8 = 7.88$ GB/s of link ceiling, and good drives land near 7 GB/s sequential. A 25 GbE link is $25/8 = 3.125$ GB/s at line rate; 10 GbE is 1.25 GB/s. Divide the checkpoint by each:

The last row is the punchline. Getting each rank's 17.64 GB shard of the 70B model from host memory onto its H100 takes about a quarter of a second. Getting the checkpoint off the media takes 20 seconds on the best local disk you can buy and nearly twelve minutes if you are pulling it from a registry over a WAN. Weight loading is a storage problem wearing a GPU costume.

Tensor-parallel ranks may iterate the same source tensors, but constructing mmap views and narrowing them does not imply physically reading every byte. Actual I/O depends on touched slices, readahead, layout, prefetch policy and cache state. Eight full physical scans are a possible unfavorable case, not a universal consequence of TP=8. Measure OS/storage counters.

The pipeline is short and every stage is a place to lose time. A checkpoint is a set of *.safetensors files plus an index JSON mapping tensor name to file. The loader opens each file, reads its header, and yields (name, tensor) pairs. The model's load_weights method translates each checkpoint name into an engine parameter name — possibly folding several checkpoint tensors into one fused parameter — and hands the tensor to that parameter's weight_loader callback. The callback narrows the tensor to this rank's slice and copies it into the already-allocated device parameter.

Two things in that diagram do the real work. The header lookup is why safetensors won; the name mapping is why adding a model is hard.

Why safetensors

Unrestricted pickle loading can execute attacker-controlled code. Modern PyTorch's restricted weights_only behavior depends on version and arguments; inspect the actual loader and never treat arbitrary checkpoint/code downloads as trusted. Safetensors uses an 8-byte little-endian header length, a JSON tensor-offset header and raw payload, avoiding executable pickle payloads and allowing mapped/partial access. This does not remove risks from model code, malformed resources or compromised dependencies.

vLLM's default path is exactly that — safe_open plus get_tensor, one tensor at a time, never holding more than one file's worth of state:

        else:
            with safe_open(st_file, framework="pt") as f:
                for name in f.keys():  # noqa: SIM118
                    if should_skip_weight(name, local_expert_ids):
                        continue
                    param = f.get_tensor(name)
                    yield name, param

SGLang's iterator is the same shape, with an explicit escape hatch for filesystems where mmap's random-access pattern is pathological:

        if disable_mmap:
            with open(st_file, "rb") as f:
                result = safetensors.torch.load(f.read())
                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():
                    yield name, f.get_tensor(name)
        if drop_cache_after_load:
            _drop_file_cache_after_load(st_file)

Network filesystems can make fault-driven loading inefficient, but one page fault is not necessarily one network round trip: caches, readahead and aggregation intervene. The following filesystem heuristic selects a strategy; compare actual I/O and memory peaks before tuning:

    fs_type = _get_fs_type(sorted_files)
    is_net_fs = fs_type in ("nfs", "nfs4", "lustre")
    total_bytes = _get_checkpoints_size_bytes(sorted_files)
    avail_bytes = _get_available_ram_bytes()
    ram_threshold_pct = 90
    fits_in_ram = total_bytes <= (ram_threshold_pct / 100.0) * avail_bytes
    fs_name = fs_type.upper() if fs_type else "unknown"

If the checkpoint is on a network FS and fits in 90% of available RAM, vLLM spawns background threads that read every file end-to-end to warm the page cache before the loader touches it — trading a sequential streaming read for the random-access pattern that would otherwise dominate. The user-facing knob is --safetensors-load-strategy, documented in vllm/config/load.py:L62-L83 with four values: lazy (mmap), eager (slurp the whole file into RAM first), prefetch, and torchao.

The name-mapping problem

A checkpoint's parameter names come from the framework that trained the model — usually HuggingFace Transformers. The engine's module tree has its own names, chosen for a different purpose: fusing GEMMs that share an input, and splitting tensors across ranks. The mapping is not one-to-one in either direction.

Case 2 is the whole difficulty. qkv_proj does not exist in the checkpoint; it is an engine invention that lets one GEMM produce Q, K and V at once. So three checkpoint tensors write into three disjoint row ranges of one engine parameter, and each of them is independently sliced by TP rank, and the K/V slicing uses a different rank arithmetic than Q because with $h_{kv}=8$ and TP=8 there is exactly one KV head per rank while there are eight query heads. Miss any of that and the model loads without error and emits fluent nonsense.

§5.1 owns why a tensor is sharded on a given axis and why the column/row split is what it is. This chapter owns the plumbing that makes rank 3 end up with rows 384–511 of k_proj and nothing else.

vLLM: a declarative mapper plus a tree walk

At a556f3f, vLLM's Llama implementation has almost no hand-written loading loop left. The mapping is a class attribute:

class LlamaModel(nn.Module, EagleModelMixin):
    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_stacked={
            # weight_name: (param_name, shard_id)
            ".q_proj": (".qkv_proj", "q"),
            ".k_proj": (".qkv_proj", "k"),
            ".v_proj": (".qkv_proj", "v"),
            ".gate_proj": (".gate_up_proj", 0),
            ".up_proj": (".gate_up_proj", 1),
        }
    )

and load_weights is three lines:

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

The mapper rewrites the name and smuggles the shard identity onto the tensor object itself:

        shard_id: ShardId | None = None
        for substr, (new_key, new_shard_id) in self.orig_to_new_stacked.items():
            if substr in key:
                key = key.replace(substr, new_key, 1)
                shard_id = new_shard_id
# ...
    def apply(
        self, weights: Iterable[tuple[str, torch.Tensor]]
    ) -> Iterable[tuple[str, torch.Tensor]]:
        for name, data in weights:
            result = self._map_name_with_shard(name)
            if result is None:
                continue
            out_name, shard_id = result
            if shard_id is not None:
                data.shard_id = shard_id
            yield out_name, data

AutoWeightsLoader.load_weights then groups the renamed stream by prefix and recurses down the module tree, so a weight named model.layers.40.self_attn.qkv_proj.weight is routed to model → layers → 40 → self_attn → qkv_proj without any string matching in the model file. When it reaches a module that defines its own load_weights, it delegates; QKVParallelLinear does, and that is where the shard id is read back off the tensor:

    def load_weights(
        self, weights: Iterable[tuple[str, torch.Tensor]]
    ) -> Iterable[str]:
        for name, loaded_weight in weights:
            shard_id = getattr(loaded_weight, "shard_id", None)
            self.validate_shard_id(shard_id)
            # Load into self if name is not an attr of self or its submodules
            param: Parameter
            if "." in name:
                submodule, _, attr = name.rpartition(".")
                param = getattr(self.get_submodule(submodule), attr, self)
            else:
                param = getattr(self, name, self)
            if param is None and name == "bias":
                continue
            param.weight_loader(param, loaded_weight, shard_id)

SGLang: the explicit loop, and a v2 behind an env var

SGLang's LlamaForCausalLM.load_weights still ships the classic table-driven loop as the default, with the tree-walking loader gated:

    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
        from sglang.srt.environ import envs

        if envs.SGLANG_ENABLE_WEIGHT_LOADER_V2.get():
            return self._load_weights_v2(weights)
        return self._legacy_load_weights(weights)

    def _legacy_load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            (".qkv_proj", ".q_proj", "q"),
            (".qkv_proj", ".k_proj", "k"),
            (".qkv_proj", ".v_proj", "v"),
            (".gate_up_proj", ".gate_proj", 0),
            (".gate_up_proj", ".up_proj", 1),
        ]

Same five-row table, different delivery. The loop that consumes it is the pattern you will find replicated across the 212 model files in python/sglang/srt/models/ (216 .py at the top level, less registry.py, utils.py and two per-family helper modules):

            for param_name, weight_name, shard_id in stacked_params_mapping:
                if weight_name not in name:
                    continue
                name = name.replace(weight_name, param_name)
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
                if name not in params_dict:
                    continue
                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
                # ...
                if name in params_dict.keys():
                    param = params_dict[name]
                    weight_loader = getattr(
                        param, "weight_loader", default_weight_loader
                    )
                    weight_loader(param, loaded_weight)
                else:
                    logger.warning(f"Parameter {name} not found in params_dict")

An unexpected source tensor and an uninitialized destination parameter are different errors. SGLang's quoted warning concerns a source name absent from its parameter map; that source may be intentionally unused. vLLM's following check instead finds expected destination parameters that were never marked loaded:

            weights_not_loaded = weights_to_load - loaded_weights
            if weights_not_loaded:
                raise ValueError(
                    "Following weights were not initialized from "
                    f"checkpoint: {weights_not_loaded}"
                )

A missing destination can have no incoming tensor to trigger an unexpected-source warning. Maintain separate expected, loaded, ignored and unexpected name sets, with intentional exceptions documented. The quoted strict check is gated by quantization and returned-name availability; other shape, format and loader-specific checks can still run.

The per-rank slice

Every path converges on a parameter's loader callback. This helper slices dimension zero; do not infer a universal row-parallel rule from its name. With stored weights [out,in], ordinary column-parallel output sharding uses dim 0 and row-parallel input sharding dim 1:

def row_parallel_weight_loader(
    param: torch.Tensor, loaded_weight: torch.Tensor
) -> None:
    """Load weights that are row-parallelized."""
    tp_rank = get_tensor_model_parallel_rank()
    shard_dim = 0 if param.dim() != 1 else None

    if shard_dim is not None:
        shard_size = param.data.shape[shard_dim]
        start_idx = tp_rank * shard_size
        loaded_weight = loaded_weight.narrow(shard_dim, start_idx, shard_size)

    return default_weight_loader(param, loaded_weight)

The fused-and-sharded case is the interesting one. Note that the destination parameter is narrowed too — the shard lands at a specific offset inside qkv_proj.weight — and that K and V use a divided rank when the model has fewer KV heads than ranks:

        assert loaded_shard_id in ["q", "k", "v"]

        # If output dim is defined, use the default loading process.
        if output_dim is not None:
            if loaded_shard_id == "q":
                shard_offset = 0
                shard_size = self.num_heads * self.head_size
            elif loaded_shard_id == "k":
                shard_offset = self.num_heads * self.head_size
                shard_size = self.num_kv_heads * self.head_size
            elif loaded_shard_id == "v":
                shard_offset = (self.num_heads + self.num_kv_heads) * self.head_size
                shard_size = self.num_kv_heads * self.v_head_size
# ...
            is_sharded_weight = getattr(param, "is_sharded_weight", False)
            param_data = param_data.narrow(output_dim, shard_offset, shard_size)
            if loaded_shard_id == "q":
                shard_rank = self.tp_rank
            else:
                shard_rank = self.tp_rank // self.num_kv_head_replicas
            start_idx = shard_rank * shard_size

            if not is_sharded_weight:
                loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)

Both narrow calls produce views, not copies. The only allocation-sized operation is the final param_data.copy_(loaded_weight), which moves exactly this rank's bytes to the device. The CPU never holds more than one checkpoint tensor at a time, and the GPU never holds a byte it does not own.

The loader zoo

vLLM's registry is a flat dict from format string to loader class (vllm/model_executor/model_loader/__init__.py:L32-L63). SGLang's is a chain of ifs in get_model_loader (python/sglang/srt/model_loader/loader.py:L4338-L4350) over a 21-member enum. What exists at these SHAs:

The shape of the difference: vLLM's zoo is about transport (how do bytes get here fastest), SGLang's is about transport plus lifecycle (how do bytes get here fastest, and how do I replace them at runtime).

Follow model.layers.40.self_attn.k_proj.weight from a Llama-3-70B checkpoint into GPU 3 at TP=8, naming every function in order.

1. get_model_loader(load_config) maps load_format="auto" to DefaultModelLoader (vllm/model_executor/model_loader/__init__.py:L48-L63).
2. BaseModelLoader.load_model enters set_default_torch_dtype(bf16), then with target_device:, then calls initialize_model — the entire module tree is constructed on the GPU with uninitialised storage, so every parameter already has its final shape and device before any byte is read (vllm/model_executor/model_loader/base_loader.py:L53-L64). Rank 3's qkv_proj.weight is allocated at [1280, 8192], not [10240, 8192].
3. DefaultModelLoader.load_weights calls model.load_weights(self.get_all_weights(...)) (vllm/model_executor/model_loader/default_loader.py:L414-L427).
4. get_all_weights → _get_weights_iterator → _prepare_weights resolves the file list, filtering against model.safetensors.index.json so a consolidated file and its shards are never both loaded.
5. safetensors_weights_iterator sniffs the filesystem, decides whether to prefetch, then safe_opens each shard and yields ("model.layers.40.self_attn.k_proj.weight", <CPU tensor 1024×8192 bf16, 16.78 MB>).
6. AutoWeightsLoader.load_weights merges in the quant-config cache-scale mapper and the rotary-embedding drop list, then calls mapper.apply(weights).
7. WeightsMapper._map_name_with_shard matches ".k_proj", rewrites the name to model.layers.40.self_attn.qkv_proj.weight, and sets data.shard_id = "k".
8. _load_module groups by prefix and recurses: model → layers → 40 → self_attn → qkv_proj. At qkv_proj the module defines load_weights, so the walk delegates.
9. QKVParallelLinear.load_weights reads shard_id back off the tensor and calls param.weight_loader(param, loaded_weight, "k").
10. QKVParallelLinear.weight_loader computes, for rank 3 with num_heads=8, num_kv_heads=1, head_size=128: shard_offset = 8*128 = 1024, shard_size = 1*128 = 128. Since num_kv_head_replicas == 1, shard_rank = 3 and start_idx = 3*128 = 384.
11. Two narrows: param_data.narrow(0, 1024, 128) selects rows 1024–1151 of the fused parameter; loaded_weight.narrow(0, 384, 128) selects rows 384–511 of the checkpoint tensor.
12. param_data.copy_(loaded_weight) copies 128*8192*2 = 2.10 MB to this GPU. The rest of the mapped tensor need not have been faulted from disk by this rank.
13. Back in BaseModelLoader.load_model: process_weights_after_loading runs per-module post-processing (weight repacking for Marlin, FP8 scale folding — see §4.3), then model.eval().
14. track_weights_loading diffs named_parameters() against the returned set and raises if anything is missing.

The two slices establish the rank's copied bytes, not its physical storage reads. Pre-sharded artifacts can simplify access and avoid some redundant reads/repacking, but their benefit must be measured with the actual iterator, storage path and cache state.

Everything above was preparation for this. Adding a model is six concrete artifacts. Read vllm/model_executor/models/llama.py — all 549 lines — before you write any of them; it is the reference implementation the other 285 model files under vllm/model_executor/models/ imitate.

Step 2 — the module tree, in the primitives

The single most common porting mistake is copying HF's modeling_*.py and leaving nn.Linear in place. Every projection must be one of ColumnParallelLinear, MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear, ReplicatedLinear, VocabParallelEmbedding, or ParallelLMHead. Choosing correctly is choosing the sharding, because each class ships the weight_loader that implements it. Llama's attention block:

        self.qkv_proj = QKVParallelLinear(
            hidden_size=hidden_size,
            head_size=self.head_dim,
            total_num_heads=self.total_num_heads,
            total_num_kv_heads=self.total_num_kv_heads,
            bias=bias,
            quant_config=quant_config,
            prefix=f"{prefix}.qkv_proj",
        )

Note that it is constructed with total head counts. The layer divides by tp_size internally, which is what makes the same model file work at TP=1 and TP=8 with no branching. The prefix is not decoration: attention layers register themselves in a global compilation and KV-cache map under their prefix, and quantization configs match per-layer ignore lists against it.

Step 3 — the mapping table

Write the WeightsMapper (vLLM) or stacked_params_mapping (SGLang) by diffing two lists. Get the checkpoint side with safetensors; get the engine side from model.named_parameters(). Anything in the first list with no home in the second is a mapping you owe. For Llama that is exactly five rows. For a model that renames mlp.w1/mlp.w3, or splits attention output into two tensors, or stores experts as a stacked 3-D tensor (§7.1), it is more.

Step 4 — registration

The two projects made opposite choices. vLLM keeps an explicit dict from HF architectures[0] to a lazily-imported module and class:

    "LlamaForCausalLM": ("llama", "LlamaForCausalLM"),
# ...
ModelRegistry = _ModelRegistry(
    {
        model_arch: _LazyRegisteredModel(
            module_name=_resolve_module_name(mod_relname),
            class_name=cls_name,
        )
        for model_arch, (mod_relname, cls_name) in _VLLM_MODELS.items()
    }
)

Laziness matters: importing 292 model modules eagerly would initialise CUDA in the parent process and break forked workers. SGLang instead scans its own package at import time and picks up any module that exports EntryClass:

    package = importlib.import_module(package_name)
    for _, name, ispkg in pkgutil.iter_modules(package.__path__, package_name + "."):
        if not ispkg:
            if name.split(".")[-1] in envs.SGLANG_DISABLED_MODEL_ARCHS.get():
                logger.debug(f"Skip loading {name} due to SGLANG_DISABLED_MODEL_ARCHS")
                continue

            try:
                module = importlib.import_module(name)
            except Exception as e:
                if strict:
                    raise
                logger.warning(f"Ignore import error when loading {name}: {e}")
                continue
            if hasattr(module, "EntryClass"):
                entry = module.EntryClass

So an SGLang contributor writes EntryClass = [MyModelForCausalLM] at the bottom of the file and is done — the trade being that a syntax error or a missing optional dependency in any model file turns into a swallowed warning and an architecture that mysteriously does not exist. vLLM's cost is a merge conflict magnet of a dict; its benefit is that a missing architecture produces a list of everything that is supported (vllm/model_executor/models/registry.py:L1167-L1169).

Step 5 — interfaces

vLLM's optional capabilities are Protocol classes with a ClassVar flag, checked at runtime with isinstance. Declaring one is a promise about your forward signature:

@runtime_checkable
class SupportsPP(Protocol):
    """The interface required for all models that support pipeline parallel."""

    supports_pp: ClassVar[Literal[True]] = True
    """
    A flag that indicates this model supports pipeline parallel.

    Note:
        There is no need to redefine this flag if this class is in the
        MRO of your model class.
    """

    make_empty_intermediate_tensors: _MakeEmptyIntermediateTensors
    """Called when PP rank > 0 for profiling purposes."""

Llama declares five of them in its class header — SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3, SupportsQuant, alongside LocalArgmaxMixin and nn.Module (vllm/model_executor/models/llama.py:L444-L452) — and SupportsLoRA is why packed_modules_mapping exists alongside the weights mapper — LoRA adapters name q_proj, so the LoRA manager needs the fused-module decomposition in reverse. Multimodal models add SupportsMultiModal, whose processing contract belongs to §7.4.

Step 6 — the tests you owe

vLLM's contributing guide is explicit that exactly one test is mandatory:

### Model loading

Include an example HuggingFace repository for your model in [tests/models/registry.py](../../../tests/models/registry.py).
This enables a unit test that loads dummy weights to ensure that the model can be initialized in vLLM.

The entry is a _HfExamplesInfo, and the extras dict is how quantized and variant checkpoints get coverage for free:

    "LlamaForCausalLM": _HfExamplesInfo(
        "meta-llama/Llama-3.2-1B-Instruct",
        extras={
            "guard": "meta-llama/Llama-Guard-3-1B",
            "hermes": "NousResearch/Hermes-3-Llama-3.1-8B",
            "fp8": "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8",
            "fp8_1b": "RedHatAI/Llama-3.2-1B-Instruct-FP8",
            "tiny": "hmellor/tiny-random-LlamaForCausalLM",
        },
    ),

Test both end-to-end logits and local mapping contracts. Direct tensor-slice equality, recognizable synthetic Q/K/V values and layerwise parity catch wrong fused offsets before generation. Text/logprob agreement is important but not the only detection method.

Three families of fix exist, and they attack different parts of the diagram.

Read fewer bytes

sharded_state is the direct attack on the 8× amplification. You pay one normal load, dump per-rank files, and every subsequent start reads only what it keeps:

class ShardedStateLoader(BaseModelLoader):
    """
    Model loader that directly loads each worker's model state dict, which
    enables a fast load path for large tensor-parallel models where each worker
    only needs to read its own shard rather than the entire checkpoint. See
    `examples/features/sharded_state/save_sharded_state_offline.py` for creating
    a sharded checkpoint.
    """

    DEFAULT_PATTERN = "model-rank-{rank}-part-{part}.safetensors"

For the 70B at TP=8 this turns 1,129 GB of logical reads into 141 GB — 8× less at worst, and it removes the dependence on having 131 GiB of spare page cache. SGLang's presharded goes further and caches the post-quantization tensors, so a run that would have spent CPU time re-deriving FP8 scales skips that too. The cost is the same in both: the dumped checkpoint is bound to one exact (TP size, PP size, quantization) tuple. Change any of them and it is garbage.

Read the bytes faster

fastsafetensors uses GPUDirect Storage to move NVMe blocks into device memory without a host bounce buffer. runai_streamer issues many concurrent range requests against object storage, which is the right shape for S3 where a single stream is latency-bound rather than bandwidth-bound. vLLM's auto-prefetch (§8.4.3) is the cheap version of the same idea for network filesystems.

Do not read the bytes at all

This is where SGLang has something vLLM does not. The weight cache daemon is a separate process that holds the fully-loaded, TP-sharded, post-quantization weights in GPU memory and hands them to engine processes as CUDA IPC handles:

"""Weight Cache Daemon — a persistent process that holds post-quantized,
TP-sharded model weights in GPU memory and serves them via CUDA IPC handles.

Each GPU runs one daemon process for its TP rank. The daemon:
1. Loads model weights from disk (full pipeline: disk → TP shard → quantize)
2. Exports every parameter/buffer as a CUDA IPC handle
3. Serves handles over a Unix socket to requesting engine processes
4. Validates CacheConfig compatibility before serving
# ...

An engine started with --weight-cache-mode client maps the daemon's tensors instead of reading a byte from disk — the whole weight-load segment of Figure 3 collapses to passing a few hundred IPC handles over a Unix socket. The loader documents a precise fallback-vs-raise contract, because a silent fallback to disk would turn a fast restart into a slow one that looks healthy:

    In client mode, disk fallback is allowed ONLY when the daemon is genuinely
    absent (its Unix socket file does not exist). Every other failure is a hard
    error rather than a silent fallback, so a broken IPC path never masquerades
    as a healthy (but slow, disk-loaded) server:

    - socket file missing            -> fall back to disk load
    - connection refused             -> raise (daemon crashed after binding)
    - CacheConfig mismatch           -> raise (do NOT disk-load on a shared GPU
                                        holding a different config's weights;
                                        also surfaces fingerprint drift bugs)
    - any protocol / transfer error  -> raise

Read the flag help before reaching for it, because the obvious usage is the wrong one:

            help="Weight cache mode. 'off': normal disk loading. "
            "'daemon': launch weight cache daemon (holds weights in GPU memory). "
            "Engine-spawned daemons are co-terminal with the engine and do NOT "
            "persist across restarts, so this alone does not speed up restart "
            "(the first start is slower). For fast recovery, run the standalone "
            "daemon (python -m sglang.srt.weight_cache.daemon) and connect with "
            "'client'. 'client': connect to existing daemon and load via IPC.",

Live weight update

The same machinery serves reinforcement-learning loops, where a trainer produces new weights every few minutes and the engine must adopt them without restarting. SGLang exposes three routes — update_weights_from_disk, update_weights_from_distributed, update_weights_from_tensor — running from the HTTP server down through the tokenizer manager, scheduler, and TP worker into the model runner. The disk route re-runs the loader against a live model, with a rollback path:

        with set_default_torch_dtype(self.model_config.dtype):
            try:
                iter = get_weight_iter(self.model_config)
            except Exception as e:
                message = f"Failed to get weights iterator: {e}."
                return False, message
            try:
                model = model_load_weights(self.get_model(), iter)
            except Exception as e:
                message = (
                    f"Failed to update weights: {e}.\nRolling back to original weights."
                )
                del iter
                gc.collect()
                iter = get_weight_iter(self.model_config)
                model_load_weights(self.get_model(), iter)
                return False, message

The tensor route is the one RL frameworks use. A trainer holding FSDP or Megatron shards calls into weight_sync/utils.py, which gathers DTensors to full tensors, serialises them as CUDA IPC handles, transposes the per-rank gather so each logical tensor carries all its shards, and posts one UpdateWeightsFromTensorReqInput:

def _preprocess_tensor_for_update_weights(tensor: torch.Tensor):
    """
    Preprocess the tensor for update weights.
    Example Use Case:
        - FSDP: we gather tensor by calling full_tensor in _preprocess_tensor_for_update_weights
        - Megatron: we do nothing here, assuming it is gathered when feed into this func

    Args:
        tensor: The tensor to be preprocessed.

    Returns:
        The full tensor if it is a DTensor, otherwise the original tensor.
    """
    if isinstance(tensor, DTensor):
        return tensor.full_tensor()
    return tensor

The iterator interface is reusable across disk, IPC and trainer-produced tensors, but load_weights is a state-mutating operation, not a pure function. Bucketed transfers amortize overhead; they do not make a multi-parameter/rank update atomic. Safe publication requires a transaction protocol described below.

The missing name: loud or silent

Three distinct failures, three different symptoms.

Unexpected checkpoint tensor. vLLM's tree walk hits a name with no module or parameter to route to:

                msg = (
                    f"There is no module or parameter named {prefix!r} "
                    f"in {self.module._get_name()}. "
                    f"The available parameters belonging to {base_prefix} "
                    f"({module._get_name()}) are: {desc_param_keys}"
                )
                raise ValueError(msg)

The error prints the available names, which is usually enough to spot the missing mapper row.

Unfilled engine parameter. Caught after the fact by track_weights_loading: "Following weights were not initialized from checkpoint: {...}". This is the one that would otherwise be silent, since an uninitialised CUDA allocation contains plausible-looking garbage.

Wrong shape. default_weight_loader's assertion fires: "Attempted to load weight (torch.Size([1024, 8192])) into parameter (torch.Size([128, 8192]))". In practice this means you narrowed with the wrong shard_size, or forgot that a GQA model's K and V are h_kv * d_h rows, not d.

Loads fine, outputs nonsense

The worst class, because nothing raises. Every shape checks out and the model emits confident garbage. It is almost always a correct-shape, wrong-content bug in a fused parameter. Bisect it like this:

1. Run at TP=1 first. Correct at TP=1 and wrong at TP=8 means shard arithmetic. Wrong at both means the mapping or the forward pass.
2. Diff against HF. Run transformers on the same checkpoint and prompt with greedy decoding and compare hidden states layer by layer. The first divergent layer names the module.
3. Check the fused offsets by hand. Assert that qkv_proj.weight[0:num_heads*head_size] equals the checkpoint's q_proj rank slice. A k/v transposition, or a gate/up swap in gate_up_proj, is invisible to every shape check because both halves have identical shape.
4. Check GQA replication. When total_num_kv_heads < tp_size, K and V are replicated and shard_rank = tp_rank // num_kv_head_replicas. Getting it wrong hands each rank the wrong KV head: grammatical output, unrelated to the prompt.
5. Turn on the loader's own logging. VLLM_LOGGING_LEVEL=DEBUG makes AutoWeightsLoader log every tensor it loads and QKVParallelLinear.load_weights log every shard. Grep for the layer you suspect.

dtype and quantization mismatches

The loader runs inside set_default_torch_dtype(model_config.dtype), so parameters are created in the engine's dtype and copy_ silently casts. An fp32 checkpoint loaded as bf16 loses mantissa bits and nobody complains — usually fine, occasionally not for models trained with unusual scaling.

Quantized checkpoints sharpen this. A quantized parameter is not one tensor but a family — packed weights, group scales, zero points, sometimes activation scales — and if the engine's quant method does not recognise the checkpoint's naming convention, some members go unmapped. vLLM's FP8 KV-scale remapper is a whole function of accumulated conventions (vllm/model_executor/model_loader/weight_utils.py:L1365-L1404), and it fails quietly:

        remapped_name = name.replace(".kv_scale", ".attn.k_scale")
        if remapped_name not in params_dict:
            logger.warning_once(
                "Found kv_scale in the checkpoint (e.g. %s), but not found the expected name in the model (e.g. %s). kv_scale is not loaded.",  #  noqa: E501
                name,
                remapped_name,
            )
            return None

A dropped KV scale means FP8 KV cache runs with a default scale, which degrades quality without erroring. This is also why track_weights_loading's strict check is disabled when model_config.quantization is set — the strictness that protects bf16 loads is exactly the strictness that a quantized checkpoint would trip on legitimately. Under quantization, you are on your own; read the warnings.

Pre-sharded checkpoints and stale configs

sharded_state and presharded dumps encode a parallelism and quantization configuration. Loading one under a different TP size gives "Could not find checkpoint files '...', only pre-sharded checkpoints are currently supported!" if the rank pattern misses, or "Missing keys {...} in loaded state!" if the file set is incomplete. Both are from vllm/model_executor/model_loader/sharded_state_loader.py:L130-L133, L161-L162. Treat a pre-sharded dump as a build artifact keyed by config, not as a checkpoint.

Before a model-scale logit test, reconstruct recognizable column/row shards and compare their forward results. For live updates: stop admission or version requests; synchronize all ranks; stage and verify complete tensors/config; preserve captured addresses or recapture; invalidate old-model KV/prefix state; atomically publish a version; then resume. Keep a validated old version or full rollback plan. Retrying a partially failed in-place loader is not proof of atomic rollback. Record model/tokenizer revisions, hashes, shard index, quantization scales and parallel layout in the artifact manifest.

import numpy as np

w = np.arange(32, dtype=float).reshape(8, 4)
x = np.array([[1., 2., 3., 4.], [-1., 0., 1., 2.]])
for tp in (1, 2):
    columns = np.split(w, tp, axis=0)
    column_output = np.concatenate([x @ shard.T for shard in columns], axis=1)
    np.testing.assert_allclose(column_output, x @ w.T)
    rows = np.split(w, tp, axis=1)
    local_x = np.split(x, tp, axis=1)
    row_output = sum(value @ shard.T for value, shard in zip(local_x, rows))
    np.testing.assert_allclose(row_output, x @ w.T)
q, k, v = [np.full((2, 4), fill) for fill in (1., 2., 3.)]
fused = np.concatenate([q, k, v], axis=0)
np.testing.assert_array_equal(fused[2:4], k)
wrong = np.concatenate([q, v, k], axis=0)
assert wrong.shape == fused.shape
assert not np.array_equal(wrong[2:4], k)
print("TP reconstruction passes; same-shape QKV swap is detected.")

All of these run on a single GPU with a small model; substitute a larger one if you have the hardware.

# Select ONE local snapshot revision, never aggregate all cached revisions.
python - /path/to/snapshots/EXACT_COMMIT <<'PY'
import json
import struct
import sys
from pathlib import Path
root = Path(sys.argv[1]).resolve(strict=True)
index = root / "model.safetensors.index.json"
if index.exists():
    names = sorted(set(json.loads(index.read_text())["weight_map"].values()))
    files = [root / name for name in names]
else:
    files = sorted(root.glob("*.safetensors"))
if not files or any(not path.is_file() for path in files):
    raise SystemExit("Missing checkpoint shards in selected revision")
if any(not path.resolve().is_relative_to(root) for path in files):
    # HF snapshots often symlink to their local blob store; inspect only trusted snapshots.
    print("Note: external/symlinked shard paths; verify the trusted snapshot manifest")
total = sum(path.stat().st_size for path in files)
tensor_count = 0
seen = set()
for path in files:
    with path.open("rb") as fh:
        prefix = fh.read(8)
        if len(prefix) != 8:
            raise ValueError("Truncated safetensors header")
        size = struct.unpack("<Q", prefix)[0]
        if size > min(path.stat().st_size - 8, 100_000_000):
            raise ValueError("Invalid or oversized header")
        header = json.loads(fh.read(size))
    tensors = {k: v for k, v in header.items() if k != "__metadata__"}
    if seen.intersection(tensors):
        raise ValueError("Duplicate tensor names across selected shards")
    seen.update(tensors)
    tensor_count += len(tensors)
print(len(files), "files;", tensor_count, "tensors;", total, "file bytes")
print("Ideal streaming seconds at 7 GB/s:", total / 7e9)
PY

# vLLM logs "Loading weights took %.2f seconds" from DefaultModelLoader.load_weights.
vllm serve meta-llama/Llama-3.2-1B-Instruct 2>&1 | grep -E "Loading weights took|Filesystem type"

# Run sequentially. Dummy vs real loading changes initialization and post-processing too.
# Their startup delta is not an exact isolated storage-time measurement.
vllm serve meta-llama/Llama-3.2-1B-Instruct --load-format dummy 2>&1 | grep "Loading weights took"

# Force the network-FS strategy on local disk and watch it get slower.
vllm serve meta-llama/Llama-3.2-1B-Instruct --safetensors-load-strategy eager

# Every mapped tensor, with its final engine name and shape.
VLLM_LOGGING_LEVEL=DEBUG vllm serve meta-llama/Llama-3.2-1B-Instruct 2>&1 \
  | grep -E "Loaded (weight|shard)" | head -40

# The same on the SGLang side, comparing the two loader paths.
python -m sglang.launch_server --model-path meta-llama/Llama-3.2-1B-Instruct
SGLANG_ENABLE_WEIGHT_LOADER_V2=1 python -m sglang.launch_server --model-path meta-llama/Llama-3.2-1B-Instruct

# In a scratch checkout of vLLM, delete the ".k_proj" row from
# LlamaModel.hf_to_vllm_mapper (vllm/model_executor/models/llama.py:L349)
# and start the server. Predict the error before you read it.
#
# Then delete only the ".up_proj" row's shard id (change 1 to 0) and start again.
# Note that this one does NOT raise. That is the whole point of §8.4.8.

• huggingface/safetensors — the format specification, including the exact header layout and the security rationale versus pickle.
• vLLM: Basic Model, Registration, and Unit Testing — the three-page contributor path this chapter grounds in source.
• Run:ai Model Streamer — the concurrent object-store reader behind --load-format runai_streamer, with the design rationale for why S3 loading is latency- rather than bandwidth-bound.
• fastsafetensors — GPUDirect Storage for safetensors, used by both engines under --load-format fastsafetensors.
• MoonshotAI/checkpoint-engine — the external parameter server SGLang integrates with in python/sglang/srt/checkpoint_engine/, aimed squarely at RL weight refresh.
• SGLang weight_cache — the daemon, the IPC loader, and the protocol; read protocol.py's CacheConfig fingerprint for how it decides two configurations are compatible.
• save_sharded_state_offline.py — the script that produces a sharded_state checkpoint, referenced directly from ShardedStateLoader's docstring.
• §5.1 for why each tensor is sharded the way it is, §5.5 for where weight load sits in the startup sequence, and §9.4 for what cold start does to an autoscaling policy.