[Model Loader][DSV41] Faster RL weight reload for MegaMoE experts

.buildkite/test_areas/models_basic.yaml

@@ -74,14 +74,15 @@ steps:
   - tests/models/test_deepseek_v41_replay_batch.py
   - tests/models/test_deepseek_v41_pcp.py
   - tests/models/test_deepseek_v41_pipeline.py
+  - tests/models/test_deepseek_v41_expert_mapping.py
   - tests/models/transformers/test_layer_registry_e2e.py
   - tests/models/test_deepseek_v41_replay_start.py
   - tests/models/test_deepseek_v4_vl_stream_weights.py
   - tests/models/test_muse_spark_rocm_loading.py
   - tests/models/test_muse_spark_logits_local.py
   - tests/models/test_muse_spark_rocm_shared_rope.py
   commands:
-    - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_pcp.py models/test_deepseek_v41_pipeline.py models/test_deepseek_v4_vl_stream_weights.py models/test_muse_spark_rocm_loading.py models/test_muse_spark_logits_local.py models/test_muse_spark_rocm_shared_rope.py
+    - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_pcp.py models/test_deepseek_v41_pipeline.py models/test_deepseek_v41_expert_mapping.py models/test_deepseek_v4_vl_stream_weights.py models/test_muse_spark_rocm_loading.py models/test_muse_spark_logits_local.py models/test_muse_spark_rocm_shared_rope.py
   mirror:
     amd:
       label: ":amd: (MI355) Basic Models (Other)"

tests/model_executor/model_loader/test_reload.py

@@ -1460,3 +1460,29 @@ def test_online_quantize_reload(
         mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0]
         add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0]
         assert add_perp < mul_perp
+
+
+class _InPlaceReloadLayer(torch.nn.Module):
+    reload_in_place = True
+
+    def __init__(self):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.zeros(2, 4), requires_grad=False)
+        self.weight.weight_loader = lambda param, loaded: param.data[0].copy_(loaded)
+
+
+def test_reload_in_place_layer_writes_live_storage(default_vllm_config):
+    """Opted-in layers keep their live tensors: no meta staging, no deferred copy."""
+    model = torch.nn.Sequential(_InPlaceReloadLayer())
+    layer = model[0]
+    live = layer.weight
+    record_metadata_for_reloading(model)
+    initialize_layerwise_reload(model)
+    assert layer.weight is live and not live.is_meta
+
+    layer.weight.weight_loader(layer.weight, torch.ones(4))
+    assert torch.equal(live[0], torch.ones(4))
+
+    finalize_layerwise_reload(model, default_vllm_config.model_config)
+    assert layer.weight is live
+    assert torch.equal(live, torch.tensor([[1.0] * 4, [0.0] * 4]))

tests/models/test_deepseek_v41_expert_mapping.py

@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+def test_expert_mapping_index_matches_linear_scan():
+    """The indexed lookup returns exactly the entries the linear scan would match."""
+    from vllm.models.deepseek_v4.nvidia.model import (
+        make_deepseek_v4_expert_params_mapping,
+    )
+    from vllm.models.deepseek_v41.nvidia.model import _expert_mapping_index
+
+    mapping = make_deepseek_v4_expert_params_mapping(384)
+    candidates = _expert_mapping_index(mapping)
+    names = [
+        f"layers.{layer}.ffn.experts.{expert}.{w}.{kind}"
+        for layer in (0, 39)
+        for expert in (0, 1, 11, 38, 111, 383)
+        for w in ("w1", "w2", "w3")
+        for kind in ("weight", "scale")
+    ] + ["layers.0.ffn.shared_experts.w1.weight", "layers.0.attn.wq_a.weight"]
+    for name in names:
+        assert candidates(name) == [m for m in mapping if m[1] in name], name
+
+    fallback = [("experts.w13_", "experts.0.gate_proj.", 0, "w1")]
+    assert _expert_mapping_index(fallback)("experts.0.gate_proj.weight") == fallback

vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py

@@ -256,6 +256,10 @@ class DeepGemmMegaMoEExperts(nn.Module):
 
     """
 
+    # Weight loaders write each checkpoint shard straight into the kernel layout
+    # and there is no post-load processing, so layerwise reload loads in place.
+    reload_in_place = True
+
     _symm_buffers: ClassVar[dict[tuple, Any]] = {}
     _synchronized_ep_groups: ClassVar[set[tuple[int, int]]] = set()
 

vllm/model_executor/model_loader/reload/layerwise.py

@@ -105,6 +105,11 @@ def initialize_layerwise_reload(model: torch.nn.Module):
         if info.can_load():
             continue
 
+        # Loaders write the final kernel layout into the live tensors and nothing
+        # post-processes them: load in place, without staging or a dry-run count.
+        if getattr(layer, "reload_in_place", False):
+            continue
+
         # Save current tensors for later copying
         info.kernel_tensors = get_layer_params_buffers(layer)
         # snapshot now: restore_layer_on_meta drops alias buffers from the live set

vllm/models/deepseek_v41/nvidia/model.py

@@ -126,6 +126,35 @@
 logger = init_logger(__name__)
 
 
+_EXPERT_WEIGHT_NAME = re.compile(r"experts\.\d+\.w[123]\.")
+
+
+def _expert_mapping_index(
+    mapping: list[tuple[str, str, int, str]],
+) -> Callable[[str], list[tuple[str, str, int, str]]]:
+    """Mapping entries whose ``weight_name`` may occur in a checkpoint name.
+
+    A linear scan of every (expert, shard) entry costs O(num_experts) per
+    expert tensor; when all weight names have the ``experts.<id>.w<k>.`` form,
+    only the entries for the occurrences of that form in the name can match.
+    """
+    index: dict[str, list[int]] = {}
+    for i, (_, weight_name, _, _) in enumerate(mapping):
+        if _EXPERT_WEIGHT_NAME.fullmatch(weight_name) is None:
+            return lambda name: mapping
+        index.setdefault(weight_name, []).append(i)
+
+    def candidates(name: str) -> list[tuple[str, str, int, str]]:
+        hits = sorted(
+            i
+            for key in _EXPERT_WEIGHT_NAME.findall(name, overlapped=True)
+            for i in index.get(key, ())
+        )
+        return [mapping[i] for i in hits]
+
+    return candidates
+
+
 class DeepseekV4MoE(DeepseekV4MoEBase):
     def __init__(
         self,
@@ -1328,6 +1357,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
 
         # Pre-compute expert mapping ONCE.
         expert_mapping = self.get_expert_mapping()
+        expert_candidates = _expert_mapping_index(expert_mapping)
 
         # Block-FP8 shared experts: pad the intermediate up to the TP-uniform
         # block count so the standard loaders below slice it evenly (trailing
@@ -1386,7 +1416,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                         and loaded_weight.dtype == torch.float8_e8m0fnu
                     ):
                         loaded_weight = loaded_weight.view(torch.uint8)
-                    for mapping in expert_mapping:
+                    for mapping in expert_candidates(name):
                         param_name, weight_name, expert_id, expert_shard_id = mapping
                         if weight_name not in name:
                             continue