[Model Loader] Direct reload: lay the checkpoint out where a relocated weight now starts

tests/model_executor/model_loader/test_direct_reload.py

@@ -289,8 +289,9 @@ class TestRefusals:
     def test_refused_begin_leaves_the_model_untouched(self):
         # proj is viewed back first, in module order; mimo's storage is short
         model = _cold_start(_checkpoint())
-        owner, (_, _, dtype, offset), attrs = model.mimo.weight._ckpt
-        model.mimo.weight._ckpt = (owner, ((64, 64), (64, 1), dtype, offset), attrs)
+        owner, (_, _, dtype, offset), attrs, storage = model.mimo.weight._ckpt
+        layout = ((64, 64), (64, 1), dtype, offset)
+        model.mimo.weight._ckpt = (owner, layout, attrs, storage)
         before = _state(model)
         with pytest.raises(RuntimeError, match="checkpoint layout needs"):
             begin_direct_reload(model)
@@ -465,6 +466,26 @@ def with_fn(self):
     assert _ptrs(_state(model)) == ptrs
 
 
+def test_reload_into_weight_moved_to_an_offset(monkeypatch):
+    def move(layer):  # DeepSeek-V4 ROCm GEMM fusion: `set_` into a shared buffer
+        w = layer.weight
+        if w.storage_offset() == 0:  # once; reloads refill it there
+            w.data = torch.cat((torch.zeros_like(w), w))[len(w) :]
+
+    init = _ToyModel.__init__
+
+    def with_move(self):
+        init(self)
+        self.proj.quant_method = _Method(move)
+
+    monkeypatch.setattr(_ToyModel, "__init__", with_move)
+    model = _cold_start(_checkpoint())
+    ptrs = _ptrs(_state(model))
+    _reload(model, _checkpoint(1.01))
+    _assert_same(_state(model), _state(_cold_start(_checkpoint(1.01))))
+    assert _ptrs(_state(model)) == ptrs
+
+
 def test_replace_parameter_keeps_the_target_checkpoint_layout():
     # post-load processing may pass another checkpoint parameter as new_tensor
     model = _cold_start(_checkpoint())

vllm/model_executor/model_loader/reload/direct.py

@@ -120,7 +120,8 @@ def record_checkpoint_layout(model: torch.nn.Module) -> None:
     for _, t in _loadable_tensors(model):
         if not t.is_meta and not hasattr(t, "_ckpt"):
             attrs = {k: t.__dict__[k] for k in _LOADER_ATTRS if k in t.__dict__}
-            t._ckpt = (model._ckpt_owner, _layout(t), attrs)
+            storage = t.untyped_storage().data_ptr()
+            t._ckpt = (model._ckpt_owner, _layout(t), attrs, storage)
     # every name, aliases included: post-load processing may replace either
     model._ckpt_param_names = frozenset(
         name
@@ -173,6 +174,8 @@ def _checkpoint_views(
             )
         shape, stride, dtype, offset = ckpt[1]
         storage = t.untyped_storage()
+        if storage.data_ptr() != ckpt[3]:  # moved, e.g. `set_` into a fused buffer
+            offset = t.storage_offset()
         span = 0 if 0 in shape else 1 + sum((n - 1) * s for n, s in zip(shape, stride))
         need = (offset + span) * dtype.itemsize
         if need > storage.nbytes():