[Bugfix][RL] Load attention sinks through weight_loader for layerwise reload

tests/model_executor/model_loader/test_reload.py

@@ -5,6 +5,7 @@
 import inspect
 import sys
 import types
+from importlib import import_module
 from types import SimpleNamespace
 from unittest.mock import Mock
 from weakref import WeakKeyDictionary, ref
@@ -39,10 +40,213 @@
 from vllm.model_executor.model_loader.weight_utils import (
     composed_weight_loader,
     default_weight_loader,
+    load_attention_sink,
 )
 from vllm.platforms import current_platform
 
 
+@pytest.mark.parametrize("tp_rank", [0, 1])
+@pytest.mark.parametrize(
+    "module_name,class_name,sink_name",
+    [
+        ("mimo_v2_mtp", "MiMoV2MTP", "attention_sink_bias"),
+        ("mimo_v2_mtp", "MiMoV2OmniMTP", "attention_sink_bias"),
+        ("laguna", "LagunaModel", "sink"),
+        ("qwen3_dflash", "DFlashQwen3Model", "attention_sink_bias"),
+    ],
+)
+def test_attention_sink_model_loading_shards_once(
+    monkeypatch, tp_rank, module_name, class_name, sink_name
+):
+    """Model loading must pass global sinks to the loader on cold load and reload."""
+    module = import_module(f"vllm.model_executor.models.{module_name}")
+    from vllm.model_executor.model_loader import weight_utils
+    from vllm.model_executor.models.utils import WeightsMapper
+
+    for target in (module, weight_utils):
+        monkeypatch.setattr(
+            target, "get_tensor_model_parallel_rank", lambda: tp_rank, raising=False
+        )
+        monkeypatch.setattr(target, "get_tensor_model_parallel_world_size", lambda: 2)
+    cls = getattr(module, class_name)
+    model = cls.__new__(cls)
+    torch.nn.Module.__init__(model)
+    model.hf_to_vllm_mapper = WeightsMapper()
+    model.model = torch.nn.Module()
+    model.model.mtp = torch.nn.Module()
+    sink = torch.nn.Parameter(torch.zeros(4), requires_grad=False)
+
+    def sink_loader(param, loaded_weight):
+        assert loaded_weight.shape == (8,), "Sink was sharded before weight_loader"
+        load_attention_sink(param, loaded_weight)
+
+    sink.weight_loader = sink_loader
+    model.model.mtp.register_parameter(sink_name, sink)
+    name = f"model.mtp.{sink_name}"
+    checkpoint = torch.arange(8, dtype=torch.float32)
+    assert model.load_weights([(name, checkpoint)]) == {name}
+    assert torch.equal(sink, checkpoint.chunk(2)[tp_rank])
+    record_metadata_for_reloading(model)
+    initialize_layerwise_reload(model)
+    model.load_weights([(name, checkpoint + 10)])
+    finalize_layerwise_reload(model, SimpleNamespace(dtype=torch.float32))
+    assert torch.equal(sink, (checkpoint + 10).chunk(2)[tp_rank])
+    assert getattr(model.model.mtp, sink_name).data_ptr() == sink.data_ptr()
+
+
+@pytest.mark.parametrize("tp_rank", [0, 1])
+@pytest.mark.parametrize(
+    "module_name,class_name,attention_name,sink_name",
+    [
+        ("deepseek_v4", "DeepseekV4Model", "attn", "attn_sink"),
+        ("deepseek_v41", "DeepseekV4Model", "attn", "attn_sink"),
+        ("hy_v4", "HYV4Model", "self_attn", "learnable_sink_param"),
+    ],
+)
+def test_attention_sink_deepseek_hy_loading_uses_parameter_loader(
+    monkeypatch, tp_rank, module_name, class_name, attention_name, sink_name
+):
+    """DeepSeek and HY-V4 pass global sinks to the parameter's current loader."""
+    module = import_module(f"vllm.models.{module_name}.nvidia.model")
+    from vllm.model_executor.model_loader import weight_utils
+
+    monkeypatch.setattr(weight_utils, "get_tensor_model_parallel_world_size", lambda: 2)
+    monkeypatch.setattr(weight_utils, "get_tensor_model_parallel_rank", lambda: tp_rank)
+    cls = getattr(module, class_name)
+    model = cls.__new__(cls)
+    torch.nn.Module.__init__(model)
+    model.config = SimpleNamespace(tie_word_embeddings=False)
+    model.quant_config = None
+    model.get_expert_mapping = lambda: []
+    attention = torch.nn.Module()
+    layer = torch.nn.Module()
+    setattr(layer, attention_name, attention)
+    model.layers = torch.nn.ModuleList([layer])
+    sink = torch.nn.Parameter(torch.zeros(6), requires_grad=False)
+    attention.register_parameter(sink_name, sink)
+    checkpoint = torch.arange(8, dtype=torch.float32)
+    name = f"layers.0.{attention_name}.{sink_name}"
+
+    for offset in (0, 10):
+        weight = checkpoint + offset
+        sink.weight_loader = Mock(side_effect=load_attention_sink)
+        assert model.load_weights([(name, weight)]) == {name}
+        sink.weight_loader.assert_called_once_with(sink, weight)
+        assert torch.equal(sink[:4], weight.chunk(2)[tp_rank])
+        assert torch.isneginf(sink[4:]).all()
+
+
+def _gpt_oss_container(monkeypatch, quant_method=None):
+    """A minimal GptOssModel whose loading paths can be driven directly."""
+    from vllm.model_executor.models import gpt_oss
+
+    group = SimpleNamespace(world_size=1, rank_in_group=0)
+    for name in ("get_dp_group", "get_ep_group", "get_pcp_group"):
+        monkeypatch.setattr(gpt_oss, name, lambda: group)
+    monkeypatch.setattr(gpt_oss, "get_tensor_model_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(gpt_oss, "get_tensor_model_parallel_rank", lambda: 0)
+    monkeypatch.setattr(gpt_oss, "is_pp_missing_parameter", lambda *args: False)
+    model = gpt_oss.GptOssModel.__new__(gpt_oss.GptOssModel)
+    torch.nn.Module.__init__(model)
+    model.config = SimpleNamespace(num_local_experts=1, intermediate_size=32)
+    if quant_method is not None:
+        model.config.quantization_config = {"quant_method": quant_method}
+    model.parallel_config = SimpleNamespace(enable_expert_parallel=True)
+    model.layers = [
+        SimpleNamespace(
+            mlp=SimpleNamespace(
+                experts=SimpleNamespace(
+                    _quant_method=SimpleNamespace(weight_dtype=None)
+                )
+            )
+        )
+    ]
+    model.get_expert_mapping = lambda: []
+    return model
+
+
+@pytest.mark.parametrize("quant_method", ["mxfp4", "quark", None])
+def test_gpt_oss_attention_sink_uses_parameter_loader(monkeypatch, quant_method):
+    """All quantization dispatch paths honor the sink's current weight loader."""
+    model = _gpt_oss_container(monkeypatch, quant_method)
+    model.sinks = torch.nn.Parameter(torch.zeros(4), requires_grad=False)
+    # A replacement loader must see the unsliced checkpoint even after reload
+    # wraps the parameter; bare model-specific copies would bypass this spy.
+    model.sinks.weight_loader = Mock(side_effect=default_weight_loader)
+    checkpoint = torch.arange(4, dtype=torch.float32)
+    assert model.load_weights([("sinks", checkpoint)]) == {"sinks"}
+    model.sinks.weight_loader.assert_called_once_with(model.sinks, checkpoint)
+    assert torch.equal(model.sinks, checkpoint)
+
+
+def test_gpt_oss_attention_sink_skips_missing_parameter(monkeypatch):
+    """A sink owned by another pipeline stage is skipped, not a KeyError."""
+    from vllm.model_executor.models import gpt_oss
+
+    model = _gpt_oss_container(monkeypatch)
+    monkeypatch.setattr(gpt_oss, "is_pp_missing_parameter", lambda *args: True)
+    # The parameter does not exist on this rank, so loading must not touch it.
+    sinks = [("sinks", torch.arange(4, dtype=torch.float32))]
+    assert model.load_weights(sinks) == set()
+
+
+@pytest.mark.parametrize("tp_rank", [0, 1])
+@pytest.mark.parametrize("padding", [0, 4])
+def test_attention_sink_layerwise_reload(monkeypatch, tp_rank, padding):
+    """Reload must update the sink's kernel storage and restore its padding."""
+    monkeypatch.setattr(
+        "vllm.model_executor.model_loader.weight_utils.get_tensor_model_parallel_rank",
+        lambda: tp_rank,
+    )
+    monkeypatch.setattr(
+        "vllm.model_executor.model_loader.weight_utils."
+        "get_tensor_model_parallel_world_size",
+        lambda: 2,
+    )
+
+    class SinkModel(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.sink = torch.nn.Parameter(
+                torch.empty(4 + padding), requires_grad=False
+            )
+            self.sink.weight_loader = load_attention_sink
+            self.proj = torch.nn.Parameter(torch.zeros(16), requires_grad=False)
+            self.proj.weight_loader = default_weight_loader
+
+        def load_weights(self, weights):
+            params = dict(self.named_parameters())
+            for name, weight in weights:
+                param = params[name]
+                param.weight_loader(param, weight)
+
+    model = SinkModel()
+    checkpoint_sink = torch.arange(8, dtype=torch.float32)
+    model.load_weights([("sink", checkpoint_sink), ("proj", torch.ones(16))])
+    kernel_sink = model.sink
+    expected = torch.cat(
+        (
+            checkpoint_sink[4 * tp_rank : 4 * (tp_rank + 1)],
+            torch.full((padding,), -torch.inf),
+        )
+    )
+    assert torch.equal(kernel_sink, expected)
+    record_metadata_for_reloading(model)
+
+    for offset in (10, 20):
+        kernel_sink.data.fill_(0)
+        initialize_layerwise_reload(model)
+        assert model.sink.is_meta
+        model.load_weights(
+            [("sink", checkpoint_sink + offset), ("proj", torch.full((16,), offset))]
+        )
+        finalize_layerwise_reload(model, SimpleNamespace(dtype=torch.float32))
+        expected[:4] = checkpoint_sink[4 * tp_rank : 4 * (tp_rank + 1)] + offset
+        assert torch.equal(kernel_sink, expected)
+        assert model.sink.data_ptr() == kernel_sink.data_ptr()
+        assert torch.equal(model.proj, torch.full((16,), offset))
+
+
 def _fp8_reload_unsupported() -> bool:
     """Whether the FP8 reload/online-quantize tests should be skipped.
 

tests/model_executor/test_weight_utils.py

@@ -9,15 +9,107 @@
 import torch
 from huggingface_hub.utils import LocalEntryNotFoundError
 
+import vllm.model_executor.model_loader.weight_utils as weight_utils
 from vllm.model_executor.model_loader.weight_utils import (
     download_weights_from_hf,
     drop_checkpoint_cache,
     initialize_single_dummy_weight,
+    load_attention_sink,
     maybe_remap_kv_scale_name,
 )
 from vllm.platforms import current_platform
 
 
+@pytest.mark.parametrize(
+    "checkpoint_heads,runtime_heads,tp_size",
+    [(8, 4, 2), (8, 8, 1), (8, 6, 2), (16, 12, 4)],
+)
+def test_attention_sink_loader_shapes(
+    monkeypatch, checkpoint_heads, runtime_heads, tp_size
+):
+    """Each rank loads its checkpoint heads and disables only padded heads."""
+    monkeypatch.setattr(
+        weight_utils, "get_tensor_model_parallel_world_size", lambda: tp_size
+    )
+    checkpoint = torch.arange(checkpoint_heads, dtype=torch.float32)
+    for tp_rank in range(tp_size):
+        monkeypatch.setattr(
+            weight_utils, "get_tensor_model_parallel_rank", lambda rank=tp_rank: rank
+        )
+        param = torch.nn.Parameter(torch.zeros(runtime_heads), requires_grad=False)
+        load_attention_sink(param, checkpoint)
+        local_heads = checkpoint.chunk(tp_size)[tp_rank]
+        expected = torch.cat(
+            (
+                local_heads,
+                torch.full((runtime_heads - local_heads.numel(),), -torch.inf),
+            )
+        )
+        assert torch.equal(param, expected)
+
+
+@pytest.mark.parametrize(
+    "checkpoint_shape,runtime_shape,error",
+    [
+        ((2, 4), (4,), RuntimeError),
+        ((8,), (2, 2), AssertionError),
+        ((7,), (4,), ValueError),
+        ((8,), (3,), AssertionError),
+    ],
+)
+def test_attention_sink_loader_rejects_invalid_shapes(
+    monkeypatch, checkpoint_shape, runtime_shape, error
+):
+    """Invalid sink shapes fail before changing the runtime parameter."""
+    monkeypatch.setattr(weight_utils, "get_tensor_model_parallel_world_size", lambda: 2)
+    monkeypatch.setattr(weight_utils, "get_tensor_model_parallel_rank", lambda: 0)
+    param = torch.nn.Parameter(torch.ones(runtime_shape), requires_grad=False)
+    with pytest.raises(error):
+        load_attention_sink(param, torch.zeros(checkpoint_shape))
+    assert torch.equal(param, torch.ones(runtime_shape))
+
+
+@pytest.mark.parametrize("param_dtype", [torch.bfloat16, torch.float32])
+def test_attention_sink_loader_keeps_param_dtype(monkeypatch, param_dtype):
+    """A checkpoint dtype different from the runtime parameter is cast, not copied."""
+    monkeypatch.setattr(weight_utils, "get_tensor_model_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(weight_utils, "get_tensor_model_parallel_rank", lambda: 0)
+    param = torch.nn.Parameter(torch.zeros(4, dtype=param_dtype), requires_grad=False)
+    load_attention_sink(param, torch.arange(4, dtype=torch.float32))
+    assert param.dtype == param_dtype
+    assert torch.equal(param.float(), torch.arange(4, dtype=torch.float32))
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a GPU")
+def test_attention_sink_loader_accepts_cpu_checkpoint(monkeypatch):
+    """Sinks are normally loaded from CPU checkpoints into device parameters."""
+    monkeypatch.setattr(weight_utils, "get_tensor_model_parallel_world_size", lambda: 2)
+    monkeypatch.setattr(weight_utils, "get_tensor_model_parallel_rank", lambda: 1)
+    param = torch.nn.Parameter(
+        torch.full((6,), 123.0, device="cuda"), requires_grad=False
+    )
+    # The test suite sets the default device, so pin the checkpoint to CPU.
+    checkpoint = torch.arange(8, dtype=torch.float32, device="cpu")
+    load_attention_sink(param, checkpoint)
+    assert param.device.type == "cuda"
+    expected = torch.arange(4, 8, dtype=torch.float32, device="cpu")
+    assert torch.equal(param[:4].cpu(), expected)
+    assert torch.isinf(param[4:]).all() and (param[4:] < 0).all()
+
+
+@pytest.mark.parametrize(
+    "num_heads,dtype", [(4, torch.float32), (8, torch.bfloat16), (6, None)]
+)
+def test_make_attention_sink_registers_loader_and_padding(num_heads, dtype):
+    """The factory is the single place that creates and registers a sink."""
+    sink = weight_utils.make_attention_sink(num_heads, dtype)
+    assert sink.shape == (num_heads,)
+    assert sink.dtype == (dtype or torch.get_default_dtype())
+    assert sink.requires_grad is False
+    assert torch.isinf(sink).all() and (sink < 0).all()
+    assert sink.weight_loader is load_attention_sink
+
+
 def test_download_weights_from_hf():
     with tempfile.TemporaryDirectory() as tmpdir:
         # assert LocalEntryNotFoundError error is thrown

tests/models/test_deepseek_v4_dspark_rocm.py

@@ -42,10 +42,6 @@ def _disable_distributed_loader_paths(monkeypatch: pytest.MonkeyPatch):
         "fused_moe_make_expert_params_mapping",
         lambda *args, **kwargs: [],
     )
-    monkeypatch.setattr(
-        dspark_module, "get_tensor_model_parallel_world_size", lambda: 1
-    )
-    monkeypatch.setattr(dspark_module, "get_tensor_model_parallel_rank", lambda: 0)
 
 
 @pytest.mark.cpu_test

tests/models/test_dspark_mla.py

@@ -258,8 +258,6 @@ def load_scale(param, weight, *args):
     draft._remap_dspark_name = lambda name: (
         dspark.DSparkDeepseekV4ForCausalLM._remap_dspark_name(draft, name)
     )
-    monkeypatch.setattr(dspark, "get_tensor_model_parallel_world_size", lambda: 4)
-    monkeypatch.setattr(dspark, "get_tensor_model_parallel_rank", lambda: 0)
     monkeypatch.setattr(
         dspark, "fused_moe_make_expert_params_mapping", lambda *a, **kw: []
     )
@@ -310,8 +308,6 @@ def capture_expert_mapping(*args, **kwargs):
     draft._remap_dspark_name = lambda name: (
         dspark.DSparkDeepseekV4ForCausalLM._remap_dspark_name(draft, name)
     )
-    monkeypatch.setattr(dspark, "get_tensor_model_parallel_world_size", lambda: 1)
-    monkeypatch.setattr(dspark, "get_tensor_model_parallel_rank", lambda: 0)
     monkeypatch.setattr(
         dspark,
         "fused_moe_make_expert_params_mapping",