diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index 66c5b49050..99b04c3cb9 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -5,6 +5,7 @@ import importlib.machinery 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.reload.utils import get_layer_tensors 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. diff --git a/tests/model_executor/test_weight_utils.py b/tests/model_executor/test_weight_utils.py index 0ecb0a1aed..3e7a0d408a 100644 --- a/tests/model_executor/test_weight_utils.py +++ b/tests/model_executor/test_weight_utils.py @@ -9,15 +9,107 @@ import pytest 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 diff --git a/tests/models/test_deepseek_v4_dspark_rocm.py b/tests/models/test_deepseek_v4_dspark_rocm.py index a5f8426e4e..34057f69fd 100644 --- a/tests/models/test_deepseek_v4_dspark_rocm.py +++ b/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 diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index 3ea010bbe5..503d03b98f 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -258,8 +258,6 @@ def test_v41_dspark_loads_linear_scales( 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 test_dsv4_dspark_fused_moe_expert_mapping_includes_redundant_experts( 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", diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index 0a834cd90c..eecdf951e5 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -33,7 +33,11 @@ from vllm.config.load import ( DEFAULT_SAFETENSORS_PREFETCH_NUM_THREADS, LoadConfig, ) -from vllm.distributed import get_tensor_model_parallel_rank, get_world_group +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + get_world_group, +) from vllm.logger import init_logger from vllm.model_executor.layers.quantization import ( QuantizationConfig, @@ -1493,6 +1497,33 @@ def row_parallel_weight_loader( LoaderFunction = Callable[[torch.Tensor, torch.Tensor], None] +def load_attention_sink(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + """Load this rank's heads of a global sink; padded heads are set to -inf.""" + tp_size = get_tensor_model_parallel_world_size() + heads = loaded_weight.numel() + if heads % tp_size: + raise ValueError( + f"Attention sink heads ({heads}) must be divisible by tp_size ({tp_size})" + ) + local_heads = heads // tp_size + start = get_tensor_model_parallel_rank() * local_heads + default_weight_loader( + param[:local_heads], loaded_weight.narrow(0, start, local_heads) + ) + param.data[local_heads:].fill_(-float("inf")) + + +def make_attention_sink( + num_heads: int, dtype: torch.dtype | None = None +) -> torch.nn.Parameter: + """Create a sink with -inf padding and register its TP weight loader.""" + sink = torch.nn.Parameter( + torch.full((num_heads,), -float("inf"), dtype=dtype), requires_grad=False + ) + sink.weight_loader = load_attention_sink + return sink + + def sharded_weight_loader(shard_axis: int) -> LoaderFunction: """Create a weight loader that shards the weights along the given axis""" diff --git a/vllm/model_executor/models/gpt_oss.py b/vllm/model_executor/models/gpt_oss.py index 90286006a6..912e3aa181 100644 --- a/vllm/model_executor/models/gpt_oss.py +++ b/vllm/model_executor/models/gpt_oss.py @@ -43,6 +43,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ) from vllm.model_executor.model_loader.weight_utils import ( default_weight_loader, + make_attention_sink, maybe_remap_kv_scale_name, remap_moe_expert_weights, ) @@ -127,9 +128,7 @@ class OAIAttention(nn.Module): tp_size = get_tensor_model_parallel_world_size() - self.sinks = torch.nn.Parameter( - torch.empty(config.num_attention_heads // tp_size, requires_grad=False) - ) + self.sinks = make_attention_sink(config.num_attention_heads // tp_size) self.q_size = self.num_attention_heads * self.head_dim // tp_size self.kv_size = self.num_key_value_heads * self.head_dim // tp_size @@ -556,8 +555,6 @@ class GptOssModel(nn.Module, EagleModelMixin): self, ep_rank_end: int, ep_rank_start: int, - heads_per_rank: int, - head_start: int, weights: Iterable[tuple[str, torch.Tensor]], stacked_params_mapping: list[tuple[str, str, str]], ) -> set[str]: @@ -720,13 +717,6 @@ class GptOssModel(nn.Module, EagleModelMixin): ) loaded_params.add(name) continue - elif "sinks" in name: - # Handle attention sinks (distributed across ranks) - param = params_dict[name] - narrow_weight = weight.narrow(0, head_start, heads_per_rank) - param.data.copy_(narrow_weight) - loaded_params.add(name) - continue for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue @@ -752,8 +742,6 @@ class GptOssModel(nn.Module, EagleModelMixin): self, ep_rank_end: int, ep_rank_start: int, - heads_per_rank: int, - head_start: int, weights: Iterable[tuple[str, torch.Tensor]], stacked_params_mapping: list[tuple[str, str, str]], ) -> set[str]: @@ -1082,14 +1070,6 @@ class GptOssModel(nn.Module, EagleModelMixin): loaded_params.add(fused_name) continue - elif "sinks" in name: - # Handle attention sinks (distributed across ranks) - param = params_dict[name] - narrow_weight = loaded_weight.narrow(0, head_start, heads_per_rank) - param.data.copy_(narrow_weight) - loaded_params.add(name) - continue - for param_name, weight_name, shard_id in stacked_params_mapping: # Skip non-stacked layers and experts (experts handled below). if weight_name not in name: @@ -1168,8 +1148,6 @@ class GptOssModel(nn.Module, EagleModelMixin): self, ep_rank_end: int, ep_rank_start: int, - heads_per_rank: int, - head_start: int, weights: Iterable[tuple[str, torch.Tensor]], stacked_params_mapping: list[tuple[str, str, str]], ) -> set[str]: @@ -1256,13 +1234,6 @@ class GptOssModel(nn.Module, EagleModelMixin): param.copy_(weight) loaded_params.add(name) continue - elif "sinks" in name: - # Handle attention sinks (distributed across ranks) - param = params_dict[name] - narrow_weight = weight.narrow(0, head_start, heads_per_rank) - param.data.copy_(narrow_weight) - loaded_params.add(name) - continue for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue @@ -1345,13 +1316,6 @@ class GptOssModel(nn.Module, EagleModelMixin): (".qkv_proj", ".v_proj", "v"), ] - tp_rank = get_tensor_model_parallel_rank() - tp_size = get_tensor_model_parallel_world_size() - - # Attention heads per rank - heads_per_rank = self.config.num_attention_heads // tp_size - head_start = tp_rank * heads_per_rank - ep_size = get_ep_group().world_size ep_rank = get_ep_group().rank_in_group num_experts = self.config.num_local_experts @@ -1383,8 +1347,6 @@ class GptOssModel(nn.Module, EagleModelMixin): return self._load_weights_mxfp4( ep_rank_end, ep_rank_start, - heads_per_rank, - head_start, weights, stacked_params_mapping, ) @@ -1392,8 +1354,6 @@ class GptOssModel(nn.Module, EagleModelMixin): return self._load_weights_quark( ep_rank_end, ep_rank_start, - heads_per_rank, - head_start, weights, stacked_params_mapping, ) @@ -1401,8 +1361,6 @@ class GptOssModel(nn.Module, EagleModelMixin): return self._load_weights_other( ep_rank_end, ep_rank_start, - heads_per_rank, - head_start, weights, stacked_params_mapping, ) diff --git a/vllm/model_executor/models/laguna.py b/vllm/model_executor/models/laguna.py index c1aaf8c9a8..55dfab6597 100644 --- a/vllm/model_executor/models/laguna.py +++ b/vllm/model_executor/models/laguna.py @@ -14,7 +14,6 @@ from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config from vllm.distributed import ( get_ep_group, get_pp_group, - get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) from vllm.logger import init_logger @@ -36,6 +35,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) +from vllm.model_executor.model_loader.weight_utils import make_attention_sink from vllm.model_executor.models.interfaces import ( EagleModelMixin, SupportsEagle3, @@ -354,9 +354,7 @@ class LagunaAttention(nn.Module): # Attention sinks (learnable per-head bias for SWA layers) sinks = None if attention_sink: - self.sink = torch.nn.Parameter( - torch.empty(self.total_num_heads // tp_size, requires_grad=False) - ) + self.sink = make_attention_sink(self.total_num_heads // tp_size) sinks = self.sink # Resolve rope params per-layer-type. ``config.rope_parameters`` is @@ -677,22 +675,6 @@ class LagunaModel(nn.Module, EagleModelMixin): return hidden_states, aux_hidden_states return hidden_states - def _slice_sink( - self, weights: Iterable[tuple[str, torch.Tensor]] - ) -> Iterable[tuple[str, torch.Tensor]]: - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() - for name, loaded_weight in weights: - # Handle attention sinks (distributed across ranks). Derive the - # per-rank slice from the parameter's own shape so per-layer - # variations in head count are handled correctly. - if "sink" in name: - heads_per_rank = loaded_weight.shape[0] // tp_size - loaded_weight = loaded_weight.narrow( - 0, tp_rank * heads_per_rank, heads_per_rank - ) - yield name, loaded_weight - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader( self, @@ -709,9 +691,7 @@ class LagunaModel(nn.Module, EagleModelMixin): "_input_scale", ], ) - return loader.load_weights( - self._slice_sink(weights), mapper=self.hf_to_vllm_mapper - ) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) class LagunaForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SupportsEagle3): diff --git a/vllm/model_executor/models/mimo_v2.py b/vllm/model_executor/models/mimo_v2.py index a6464c7c84..c184655676 100644 --- a/vllm/model_executor/models/mimo_v2.py +++ b/vllm/model_executor/models/mimo_v2.py @@ -47,6 +47,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ) from vllm.model_executor.model_loader.weight_utils import ( default_weight_loader, + make_attention_sink, maybe_remap_kv_scale_name, ) from vllm.model_executor.models.utils import sequence_parallel_chunk @@ -290,9 +291,7 @@ class MiMoV2Attention(nn.Module): ) self.attention_sink_bias = ( - torch.nn.Parameter(torch.empty(self.num_heads), requires_grad=False) - if add_swa_attention_sink_bias - else None + make_attention_sink(self.num_heads) if add_swa_attention_sink_bias else None ) sliding_window = sliding_window_size if sliding_window_size > -1 else None @@ -848,18 +847,9 @@ class MiMoV2Model(nn.Module, EagleModelMixin): param = params_dict[name] - if "attention_sink_bias" in name: - total_heads = loaded_weight.shape[0] - heads_per_rank = total_heads // tp_size - head_start = tp_rank * heads_per_rank - narrow_weight = loaded_weight.narrow(0, head_start, heads_per_rank) - - param.data.copy_(narrow_weight) - loaded_params.add(name) - else: - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) return loaded_params diff --git a/vllm/model_executor/models/mimo_v2_mtp.py b/vllm/model_executor/models/mimo_v2_mtp.py index fe6d64552d..d4b7556c3f 100644 --- a/vllm/model_executor/models/mimo_v2_mtp.py +++ b/vllm/model_executor/models/mimo_v2_mtp.py @@ -377,14 +377,6 @@ class MiMoV2MTP(nn.Module): continue param = params_dict[name] - # attention_sink_bias is head-parallel; slice by tp - if "attention_sink_bias" in name: - total_heads = loaded_weight.shape[0] - heads_per_rank = total_heads // tp_size - loaded_weight = loaded_weight.narrow( - 0, tp_rank * heads_per_rank, heads_per_rank - ) - direct_weight_loader = getattr( param, "weight_loader", default_weight_loader ) diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 44b7865daa..485ea952fc 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -14,7 +14,6 @@ from vllm import _custom_ops as ops from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config from vllm.distributed import ( - get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) from vllm.logger import init_logger @@ -33,6 +32,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) +from vllm.model_executor.model_loader.weight_utils import make_attention_sink from vllm.multimodal.inputs import NestedTensors from vllm.platforms import current_platform from vllm.transformers_utils.config import set_default_rope_theta @@ -251,9 +251,7 @@ class DFlashQwen3Attention(nn.Module): ) self.attention_sink_bias = ( - torch.nn.Parameter(torch.empty(self.num_heads), requires_grad=False) - if add_swa_attention_sink_bias - else None + make_attention_sink(self.num_heads) if add_swa_attention_sink_bias else None ) self.sliding_window = sliding_window @@ -746,26 +744,9 @@ class DFlashQwen3Model(nn.Module): hidden_states, _ = self.norm(hidden_states, residual) return hidden_states - def _preprocess( - self, weights: Iterable[tuple[str, torch.Tensor]] - ) -> Iterable[tuple[str, torch.Tensor]]: - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() - for name, loaded_weight in weights: - if "attention_sink_bias" in name: - # Sink bias is per-head; shard it across TP ranks like the - # attention heads themselves. - heads_per_rank = loaded_weight.shape[0] // tp_size - loaded_weight = loaded_weight.narrow( - 0, tp_rank * heads_per_rank, heads_per_rank - ) - yield name, loaded_weight - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) - return loader.load_weights( - self._preprocess(weights), mapper=self.hf_to_vllm_mapper - ) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) class DFlashQwen3ForCausalLM(Qwen3ForCausalLM): diff --git a/vllm/models/deepseek_v4/amd/dspark.py b/vllm/models/deepseek_v4/amd/dspark.py index e62431d7b6..a316b6927e 100644 --- a/vllm/models/deepseek_v4/amd/dspark.py +++ b/vllm/models/deepseek_v4/amd/dspark.py @@ -26,10 +26,6 @@ import torch import torch.nn as nn from vllm.config import VllmConfig, get_current_vllm_config -from vllm.distributed import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, @@ -413,12 +409,6 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): loaded_params: set[str] = set() loaded_confidence_head = False - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() - n_local_head = self.config.num_attention_heads // tp_size - head_start = n_local_head * tp_rank - head_end = n_local_head * (tp_rank + 1) - for name, loaded_weight in weights: mapped = self._remap_dspark_name(name) if mapped is None: @@ -474,11 +464,6 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): loaded_params.add(name) break else: - if "attn_sink" in name: - narrow = loaded_weight[head_start:head_end] - params_dict[name][: narrow.shape[0]].copy_(narrow) - loaded_params.add(name) - continue if ".shared_experts.w2" in name: name = name.replace( ".shared_experts.w2", ".shared_experts.down_proj" diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 51d22b75e8..76db8d132f 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -1265,14 +1265,6 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): return inv_name return name - # TP for attention - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() - n_head = self.config.num_attention_heads - n_local_head = n_head // tp_size - head_rank_start = n_local_head * tp_rank - head_rank_end = n_local_head * (tp_rank + 1) - # Pre-compute expert mapping ONCE. expert_mapping = self.get_expert_mapping() @@ -1362,14 +1354,6 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): break loaded_params.add(name_mapped) continue - elif "attn_sink" in name: - if is_pp_missing_parameter(name, self): - continue - narrow_weight = loaded_weight[head_rank_start:head_rank_end] - n = narrow_weight.shape[0] - params_dict[name][:n].copy_(narrow_weight) - loaded_params.add(name) - continue else: if is_pp_missing_parameter(name, self): continue diff --git a/vllm/models/deepseek_v4/amd/mtp.py b/vllm/models/deepseek_v4/amd/mtp.py index 3746d6e387..6ddac1bc05 100644 --- a/vllm/models/deepseek_v4/amd/mtp.py +++ b/vllm/models/deepseek_v4/amd/mtp.py @@ -19,10 +19,6 @@ import torch import torch.nn as nn from vllm.config import VllmConfig -from vllm.distributed import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import fused_moe_make_expert_params_mapping from vllm.model_executor.layers.layernorm import RMSNorm @@ -384,14 +380,6 @@ class DeepSeekV4MTP(nn.Module): # KeyError downstream rather than be silently rewritten. return name - # TP for attention - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() - n_head = self.config.num_attention_heads - n_local_head = n_head // tp_size - head_rank_start = n_local_head * tp_rank - head_rank_end = n_local_head * (tp_rank + 1) - # Pre-compute expert mapping ONCE. first_layer = next(iter(self.model.layers.values())) expert_mapping = fused_moe_make_expert_params_mapping( @@ -489,12 +477,6 @@ class DeepSeekV4MTP(nn.Module): loaded_params.add(name_mapped) break continue - elif "attn_sink" in name: - narrow_weight = loaded_weight[head_rank_start:head_rank_end] - n = narrow_weight.shape[0] - params_dict[name][:n].copy_(narrow_weight) - loaded_params.add(name) - continue else: if ".shared_experts.w2" in name: name = name.replace( diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 4f756538fb..2542897ef8 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -20,6 +20,7 @@ from vllm.model_executor.layers.linear import ( RowParallelLinear, ) from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer +from vllm.model_executor.model_loader.weight_utils import make_attention_sink from vllm.models.common.ops.fused_qk_rmsnorm import ( _FUSED_Q_KV_RMSNORM_KERNEL, fused_q_kv_rmsnorm, @@ -240,10 +241,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): self.padded_heads = self.get_padded_num_q_heads(self.n_local_heads) # Sink padded to the same head count, initialized to -inf (no sink # effect). Weight loading fills the first n_local_heads slots. - self.attn_sink = nn.Parameter( - torch.full((self.padded_heads,), -float("inf"), dtype=torch.float32), - requires_grad=False, - ) + self.attn_sink = make_attention_sink(self.padded_heads, torch.float32) self.fused_wqa_wkv = MergedColumnParallelLinear( self.hidden_size, diff --git a/vllm/models/deepseek_v4/cpu/model.py b/vllm/models/deepseek_v4/cpu/model.py index 6e88d21c27..af66c480fc 100644 --- a/vllm/models/deepseek_v4/cpu/model.py +++ b/vllm/models/deepseek_v4/cpu/model.py @@ -663,14 +663,6 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): return inv_name return name - # TP for attention - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() - n_head = self.config.num_attention_heads - n_local_head = n_head // tp_size - head_rank_start = n_local_head * tp_rank - head_rank_end = n_local_head * (tp_rank + 1) - # Pre-compute expert mapping ONCE. expert_mapping = self.get_expert_mapping() @@ -755,14 +747,6 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): break loaded_params.add(name_mapped) continue - elif "attn_sink" in name: - if is_pp_missing_parameter(name, self): - continue - narrow_weight = loaded_weight[head_rank_start:head_rank_end] - n = narrow_weight.shape[0] - params_dict[name][:n].copy_(narrow_weight) - loaded_params.add(name) - continue else: if is_pp_missing_parameter(name, self): continue diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index db47b7cf11..dda1026a53 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -18,10 +18,6 @@ import torch.nn as nn import vllm.envs as envs from vllm.config import VllmConfig, get_current_vllm_config from vllm.config.kernel import NATIVE_MEGA_MOE_BACKENDS -from vllm.distributed import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger from vllm.model_executor.kernels.mhc.tilelang import ( @@ -484,12 +480,6 @@ class DSparkDeepseekV4ForCausalLM(nn.Module, DeepseekV4MixtureOfExperts): loaded_params: set[str] = set() loaded_confidence_head = False - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() - n_local_head = self.config.num_attention_heads // tp_size - head_start = n_local_head * tp_rank - head_end = n_local_head * (tp_rank + 1) - weights = _duplicate_context_wkv_weights(weights, len(self.model.layers)) for name, loaded_weight in weights: mapped = self._remap_dspark_name(name) @@ -565,11 +555,6 @@ class DSparkDeepseekV4ForCausalLM(nn.Module, DeepseekV4MixtureOfExperts): loaded_params.add(name) break else: - if "attn_sink" in name: - narrow = loaded_weight[head_start:head_end] - params_dict[name][: narrow.shape[0]].copy_(narrow) - loaded_params.add(name) - continue if name.endswith(".ffn.gate.bias"): name = name.replace( ".ffn.gate.bias", ".ffn.gate.e_score_correction_bias" diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 6dd49f9409..275dc33491 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -14,7 +14,6 @@ from vllm.config.kernel import MEGA_MOE_BACKENDS, NATIVE_MEGA_MOE_BACKENDS from vllm.distributed import ( get_ep_group, get_pp_group, - get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) from vllm.distributed.eplb.eplb_state import EplbLayerState @@ -1399,14 +1398,6 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() - # TP for attention - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() - n_head = self.config.num_attention_heads - n_local_head = n_head // tp_size - head_rank_start = n_local_head * tp_rank - head_rank_end = n_local_head * (tp_rank + 1) - # Pre-compute expert mapping ONCE. expert_mapping = self.get_expert_mapping() @@ -1490,14 +1481,6 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): break loaded_params.add(name_mapped) continue - elif "attn_sink" in name: - if is_pp_missing_parameter(name, self): - continue - narrow_weight = loaded_weight[head_rank_start:head_rank_end] - n = narrow_weight.shape[0] - params_dict[name][:n].copy_(narrow_weight) - loaded_params.add(name) - continue else: if is_pp_missing_parameter(name, self): continue diff --git a/vllm/models/deepseek_v4/nvidia/mtp.py b/vllm/models/deepseek_v4/nvidia/mtp.py index 0956ab053f..838dd66945 100644 --- a/vllm/models/deepseek_v4/nvidia/mtp.py +++ b/vllm/models/deepseek_v4/nvidia/mtp.py @@ -21,10 +21,6 @@ import torch.nn as nn import vllm.envs as envs from vllm.config import VllmConfig from vllm.config.kernel import NATIVE_MEGA_MOE_BACKENDS -from vllm.distributed import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger from vllm.model_executor.kernels.mhc.tilelang import ( @@ -386,14 +382,6 @@ class DeepSeekV4MTP(nn.Module): params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() - # TP for attention - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() - n_head = self.config.num_attention_heads - n_local_head = n_head // tp_size - head_rank_start = n_local_head * tp_rank - head_rank_end = n_local_head * (tp_rank + 1) - # Pre-compute expert mapping ONCE. first_layer = next(iter(self.model.layers.values())) if first_layer.mtp_block.ffn.use_native_mega_moe: @@ -497,12 +485,6 @@ class DeepSeekV4MTP(nn.Module): loaded_params.add(name_mapped) break continue - elif "attn_sink" in name: - narrow_weight = loaded_weight[head_rank_start:head_rank_end] - n = narrow_weight.shape[0] - params_dict[name][:n].copy_(narrow_weight) - loaded_params.add(name) - continue else: if name.endswith(".ffn.gate.bias"): # ``e_score_correction_bias`` lives on the gate diff --git a/vllm/models/deepseek_v4/xpu/dspark.py b/vllm/models/deepseek_v4/xpu/dspark.py index 63333a9c0a..41ec1b4a16 100644 --- a/vllm/models/deepseek_v4/xpu/dspark.py +++ b/vllm/models/deepseek_v4/xpu/dspark.py @@ -15,10 +15,6 @@ import torch import torch.nn as nn from vllm.config import VllmConfig, get_current_vllm_config -from vllm.distributed import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, @@ -322,12 +318,6 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() - n_local_head = self.config.num_attention_heads // tp_size - head_start = n_local_head * tp_rank - head_end = n_local_head * (tp_rank + 1) - for name, loaded_weight in weights: mapped = self._remap_dspark_name(name) if mapped is None: @@ -383,13 +373,6 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): loaded_params.add(name) break else: - if "attn_sink" in name: - if name not in params_dict: - continue - narrow = loaded_weight[head_start:head_end] - params_dict[name][: narrow.shape[0]].copy_(narrow) - loaded_params.add(name) - continue if ".shared_experts.w2" in name: name = name.replace( ".shared_experts.w2", ".shared_experts.down_proj" diff --git a/vllm/models/deepseek_v4/xpu/model.py b/vllm/models/deepseek_v4/xpu/model.py index f3e468ebb1..92093a15a2 100644 --- a/vllm/models/deepseek_v4/xpu/model.py +++ b/vllm/models/deepseek_v4/xpu/model.py @@ -1149,14 +1149,6 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() - # TP for attention - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() - n_head = self.config.num_attention_heads - n_local_head = n_head // tp_size - head_rank_start = n_local_head * tp_rank - head_rank_end = n_local_head * (tp_rank + 1) - # Pre-compute expert mapping ONCE. expert_mapping = self.get_expert_mapping() @@ -1219,14 +1211,6 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): break loaded_params.add(name_mapped) continue - elif "attn_sink" in name: - if is_pp_missing_parameter(name, self): - continue - narrow_weight = loaded_weight[head_rank_start:head_rank_end] - n = narrow_weight.shape[0] - params_dict[name][:n].copy_(narrow_weight) - loaded_params.add(name) - continue else: if is_pp_missing_parameter(name, self): continue diff --git a/vllm/models/deepseek_v4/xpu/mtp.py b/vllm/models/deepseek_v4/xpu/mtp.py index a19d0d5782..dfbb79722a 100644 --- a/vllm/models/deepseek_v4/xpu/mtp.py +++ b/vllm/models/deepseek_v4/xpu/mtp.py @@ -19,10 +19,6 @@ import torch import torch.nn as nn from vllm.config import VllmConfig -from vllm.distributed import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import fused_moe_make_expert_params_mapping from vllm.model_executor.layers.layernorm import RMSNorm @@ -335,14 +331,6 @@ class DeepSeekV4MTP(nn.Module): params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() - # TP for attention - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() - n_head = self.config.num_attention_heads - n_local_head = n_head // tp_size - head_rank_start = n_local_head * tp_rank - head_rank_end = n_local_head * (tp_rank + 1) - # Pre-compute expert mapping ONCE. first_layer = next(iter(self.model.layers.values())) if first_layer.mtp_block.ffn.use_mega_moe: @@ -441,12 +429,6 @@ class DeepSeekV4MTP(nn.Module): loaded_params.add(name_mapped) break continue - elif "attn_sink" in name: - narrow_weight = loaded_weight[head_rank_start:head_rank_end] - n = narrow_weight.shape[0] - params_dict[name][:n].copy_(narrow_weight) - loaded_params.add(name) - continue else: if ".shared_experts.w2" in name: name = name.replace( diff --git a/vllm/models/deepseek_v41/amd/dspark.py b/vllm/models/deepseek_v41/amd/dspark.py index 78f07b6b11..14c6c97257 100644 --- a/vllm/models/deepseek_v41/amd/dspark.py +++ b/vllm/models/deepseek_v41/amd/dspark.py @@ -18,10 +18,6 @@ import torch.nn as nn import vllm.envs as envs from vllm.config import VllmConfig, get_current_vllm_config -from vllm.distributed import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger from vllm.model_executor.kernels.mhc.torch import mhc_post_torch @@ -420,12 +416,6 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): loaded_params: set[str] = set() loaded_confidence_head = False - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() - n_local_head = self.config.num_attention_heads // tp_size - head_start = n_local_head * tp_rank - head_end = n_local_head * (tp_rank + 1) - for name, loaded_weight in weights: mapped = self._remap_dspark_name(name) if mapped is None: @@ -487,11 +477,6 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): loaded_params.add(name) break else: - if "attn_sink" in name: - narrow = loaded_weight[head_start:head_end] - params_dict[name][: narrow.shape[0]].copy_(narrow) - loaded_params.add(name) - continue if name.endswith(".ffn.gate.bias"): name = name.replace( ".ffn.gate.bias", ".ffn.gate.e_score_correction_bias" diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index a52a408c74..639bea27ad 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -12,7 +12,6 @@ import vllm.envs as envs from vllm.config import VllmConfig from vllm.distributed import ( get_pp_group, - get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, ) @@ -985,14 +984,6 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): return inv_name return name - # TP for attention - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() - n_head = self.config.num_attention_heads - n_local_head = n_head // tp_size - head_rank_start = n_local_head * tp_rank - head_rank_end = n_local_head * (tp_rank + 1) - # Pre-compute expert mapping ONCE. expert_mapping = self.get_expert_mapping() @@ -1073,14 +1064,6 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): break loaded_params.add(name_mapped) continue - elif "attn_sink" in name: - if is_pp_missing_parameter(name, self): - continue - narrow_weight = loaded_weight[head_rank_start:head_rank_end] - n = narrow_weight.shape[0] - params_dict[name][:n].copy_(narrow_weight) - loaded_params.add(name) - continue else: if is_pp_missing_parameter(name, self): continue diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index 4d460d89b8..7009b466a2 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -25,6 +25,7 @@ from vllm.model_executor.layers.linear import ( ) from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer from vllm.model_executor.layers.sparse_mqa_indexer import SparseMQAIndexer +from vllm.model_executor.model_loader.weight_utils import make_attention_sink from vllm.models.common.ops import fused_q_kv_rmsnorm from vllm.models.common.ops.sequence_parallel import sp_reduce_scatter from vllm.models.deepseek_v41.common.ops import ( @@ -356,10 +357,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): self.padded_heads = self.get_padded_num_q_heads(self.n_local_heads) # Sink padded to the same head count, initialized to -inf (no sink # effect). Weight loading fills the first n_local_heads slots. - self.attn_sink = nn.Parameter( - torch.full((self.padded_heads,), -float("inf"), dtype=torch.float32), - requires_grad=False, - ) + self.attn_sink = make_attention_sink(self.padded_heads, torch.float32) self.fused_wqa_wkv = MergedColumnParallelLinear( self.hidden_size, diff --git a/vllm/models/deepseek_v41/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py index 0e9c43d64a..c1310263a8 100644 --- a/vllm/models/deepseek_v41/nvidia/dspark.py +++ b/vllm/models/deepseek_v41/nvidia/dspark.py @@ -19,10 +19,6 @@ import torch.nn as nn import vllm.envs as envs from vllm.config import VllmConfig, get_current_vllm_config from vllm.config.kernel import NATIVE_MEGA_MOE_BACKENDS -from vllm.distributed import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger from vllm.model_executor.kernels.mhc.tilelang import ( @@ -409,12 +405,6 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): loaded_params: set[str] = set() loaded_confidence_head = False - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() - n_local_head = self.config.num_attention_heads // tp_size - head_start = n_local_head * tp_rank - head_end = n_local_head * (tp_rank + 1) - for name, loaded_weight in weights: mapped = self._remap_dspark_name(name) if mapped is None: @@ -483,11 +473,6 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): loaded_params.add(name) break else: - if "attn_sink" in name: - narrow = loaded_weight[head_start:head_end] - params_dict[name][: narrow.shape[0]].copy_(narrow) - loaded_params.add(name) - continue if name.endswith(".ffn.gate.bias"): name = name.replace( ".ffn.gate.bias", ".ffn.gate.e_score_correction_bias" diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index bc516c8de7..3757809dbc 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -15,7 +15,6 @@ from vllm.config.kernel import MEGA_MOE_BACKENDS, NATIVE_MEGA_MOE_BACKENDS from vllm.distributed import ( get_engram_dp_size, get_pp_group, - get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, ) @@ -1251,14 +1250,6 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() - # TP for attention - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() - n_head = self.config.num_attention_heads - n_local_head = n_head // tp_size - head_rank_start = n_local_head * tp_rank - head_rank_end = n_local_head * (tp_rank + 1) - # Pre-compute expert mapping ONCE. expert_mapping = self.get_expert_mapping() @@ -1346,14 +1337,6 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): break loaded_params.add(name_mapped) continue - elif "attn_sink" in name: - if is_pp_missing_parameter(name, self): - continue - narrow_weight = loaded_weight[head_rank_start:head_rank_end] - n = narrow_weight.shape[0] - params_dict[name][:n].copy_(narrow_weight) - loaded_params.add(name) - continue else: if is_pp_missing_parameter(name, self): continue diff --git a/vllm/models/hy_v4/amd/model.py b/vllm/models/hy_v4/amd/model.py index 772e4e5b74..50925b060c 100644 --- a/vllm/models/hy_v4/amd/model.py +++ b/vllm/models/hy_v4/amd/model.py @@ -36,8 +36,6 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig from vllm.distributed import ( get_pp_group, - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, ) from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import fused_moe_make_expert_params_mapping @@ -435,12 +433,6 @@ class HYV4Model(nn.Module): split_expert_params_mapping = self.get_expert_mapping() loaded_params: set[str] = set() - # Sink weights are sharded like the q/k/v linears. - sink_tp_size = get_tensor_model_parallel_world_size() - sink_tp_rank = get_tensor_model_parallel_rank() - n_local_head = self.config.num_attention_heads // sink_tp_size - head_rank_start = n_local_head * sink_tp_rank - head_rank_end = n_local_head * (sink_tp_rank + 1) base_layer = ( "base_layer." if any(".base_layer." in name for name in params_dict) else "" ) @@ -572,48 +564,36 @@ class HYV4Model(nn.Module): if loaded_expert_param_names: loaded_params.update(loaded_expert_param_names) continue - if "learnable_sink_param" in name: - if is_pp_missing_parameter(name, self): - continue - narrow_weight = loaded_weight[head_rank_start:head_rank_end] - n = narrow_weight.shape[0] - with torch.no_grad(): - params_dict[name][:n].copy_(narrow_weight) - else: - if is_expert_weight: - # An expert weight that is not mapped to this rank. - continue - if "gate.e_score_correction_bias" in name: - name = name.replace( - "gate.e_score_correction_bias", "expert_bias" - ) - if "router.gate." in name: - name = name.replace("router.", "") - if "hc_fn" in name: - name = name.replace("hc_fn", "hc_fn.weight") - if "hc_head_fn" in name: - name = name.replace("hc_head_fn", "hc_head_fn.weight") + if is_expert_weight: + # An expert weight that is not mapped to this rank. + continue + if "gate.e_score_correction_bias" in name: + name = name.replace("gate.e_score_correction_bias", "expert_bias") + if "router.gate." in name: + name = name.replace("router.", "") + if "hc_fn" in name: + name = name.replace("hc_fn", "hc_fn.weight") + if "hc_head_fn" in name: + name = name.replace("hc_head_fn", "hc_head_fn.weight") - if is_pp_missing_parameter(name, self): - continue + if is_pp_missing_parameter(name, self): + continue - if name not in params_dict: - if _should_skip_missing_param(name): - continue - logger.warning_once( - "Skipping unknown checkpoint weight: %s", - name, - ) + if name not in params_dict: + if _should_skip_missing_param(name): continue - - param = params_dict[name] - weight_loader = getattr( - param, "weight_loader", default_weight_loader + logger.warning_once( + "Skipping unknown checkpoint weight: %s", + name, ) - if stacked_shard_id is None: - weight_loader(param, loaded_weight) - else: - weight_loader(param, loaded_weight, stacked_shard_id) + continue + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + if stacked_shard_id is None: + weight_loader(param, loaded_weight) + else: + weight_loader(param, loaded_weight, stacked_shard_id) loaded_params.add(name) return loaded_params diff --git a/vllm/models/hy_v4/nvidia/attention.py b/vllm/models/hy_v4/nvidia/attention.py index 48d769a856..d10d3afe90 100644 --- a/vllm/models/hy_v4/nvidia/attention.py +++ b/vllm/models/hy_v4/nvidia/attention.py @@ -40,6 +40,7 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import ( ) from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer +from vllm.model_executor.model_loader.weight_utils import make_attention_sink from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache from vllm.platforms import current_platform from vllm.v1.attention.backend import AttentionBackend, AttentionType @@ -493,13 +494,11 @@ class HYV4MLAAttention(nn.Module): if self.learnable_sink: sink_backend = self._resolve_sink_backend(kv_cache_dtype) enable_sink = sink_backend is not None - self.learnable_sink_param = nn.Parameter( - torch.empty( - self.num_local_heads, - # The kernels require fp32 sinks; the disabled path keeps - # the checkpoint dtype since the value is never consumed. - dtype=torch.float32 if enable_sink else torch.bfloat16, - ) + # The kernels require fp32 sinks; the disabled path keeps the + # checkpoint dtype since the value is never consumed. + self.learnable_sink_param = make_attention_sink( + self.num_local_heads, + torch.float32 if enable_sink else torch.bfloat16, ) if enable_sink: sinks = self.learnable_sink_param diff --git a/vllm/models/hy_v4/nvidia/model.py b/vllm/models/hy_v4/nvidia/model.py index e9313989fa..6e91ec2422 100644 --- a/vllm/models/hy_v4/nvidia/model.py +++ b/vllm/models/hy_v4/nvidia/model.py @@ -35,8 +35,6 @@ from transformers import PreTrainedConfig from vllm.config import CacheConfig, VllmConfig from vllm.distributed import ( get_pp_group, - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, ) from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import fused_moe_make_expert_params_mapping @@ -433,12 +431,6 @@ class HYV4Model(nn.Module): split_expert_params_mapping = self.get_expert_mapping() loaded_params: set[str] = set() - # Sink weights are sharded like the q/k/v linears. - sink_tp_size = get_tensor_model_parallel_world_size() - sink_tp_rank = get_tensor_model_parallel_rank() - n_local_head = self.config.num_attention_heads // sink_tp_size - head_rank_start = n_local_head * sink_tp_rank - head_rank_end = n_local_head * (sink_tp_rank + 1) base_layer = ( "base_layer." if any(".base_layer." in name for name in params_dict) else "" ) @@ -570,48 +562,36 @@ class HYV4Model(nn.Module): if loaded_expert_param_names: loaded_params.update(loaded_expert_param_names) continue - if "learnable_sink_param" in name: - if is_pp_missing_parameter(name, self): - continue - narrow_weight = loaded_weight[head_rank_start:head_rank_end] - n = narrow_weight.shape[0] - with torch.no_grad(): - params_dict[name][:n].copy_(narrow_weight) - else: - if is_expert_weight: - # An expert weight that is not mapped to this rank. - continue - if "gate.e_score_correction_bias" in name: - name = name.replace( - "gate.e_score_correction_bias", "expert_bias" - ) - if "router.gate." in name: - name = name.replace("router.", "") - if "hc_fn" in name: - name = name.replace("hc_fn", "hc_fn.weight") - if "hc_head_fn" in name: - name = name.replace("hc_head_fn", "hc_head_fn.weight") + if is_expert_weight: + # An expert weight that is not mapped to this rank. + continue + if "gate.e_score_correction_bias" in name: + name = name.replace("gate.e_score_correction_bias", "expert_bias") + if "router.gate." in name: + name = name.replace("router.", "") + if "hc_fn" in name: + name = name.replace("hc_fn", "hc_fn.weight") + if "hc_head_fn" in name: + name = name.replace("hc_head_fn", "hc_head_fn.weight") - if is_pp_missing_parameter(name, self): - continue + if is_pp_missing_parameter(name, self): + continue - if name not in params_dict: - if _should_skip_missing_param(name): - continue - logger.warning_once( - "Skipping unknown checkpoint weight: %s", - name, - ) + if name not in params_dict: + if _should_skip_missing_param(name): continue - - param = params_dict[name] - weight_loader = getattr( - param, "weight_loader", default_weight_loader + logger.warning_once( + "Skipping unknown checkpoint weight: %s", + name, ) - if stacked_shard_id is None: - weight_loader(param, loaded_weight) - else: - weight_loader(param, loaded_weight, stacked_shard_id) + continue + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + if stacked_shard_id is None: + weight_loader(param, loaded_weight) + else: + weight_loader(param, loaded_weight, stacked_shard_id) loaded_params.add(name) return loaded_params diff --git a/vllm/models/hy_v4/nvidia/mtp.py b/vllm/models/hy_v4/nvidia/mtp.py index 8683777cf4..6647cda8a5 100644 --- a/vllm/models/hy_v4/nvidia/mtp.py +++ b/vllm/models/hy_v4/nvidia/mtp.py @@ -11,10 +11,6 @@ from torch import nn from transformers import PreTrainedConfig from vllm.config import CacheConfig, ModelConfig, VllmConfig -from vllm.distributed import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import fused_moe_make_expert_params_mapping from vllm.model_executor.layers.layernorm import RMSNorm @@ -744,11 +740,6 @@ class HYV4MTP(nn.Module): } num_experts = getattr(self.config, "n_routed_experts", 0) - sink_tp_size = get_tensor_model_parallel_world_size() - sink_tp_rank = get_tensor_model_parallel_rank() - n_local_head = self.config.num_attention_heads // sink_tp_size - head_rank_start = n_local_head * sink_tp_rank - head_rank_end = n_local_head * (sink_tp_rank + 1) # Routed-expert weights come in two checkpoint layouts: # split: mlp.experts..gate_proj.weight @@ -883,15 +874,6 @@ class HYV4MTP(nn.Module): if is_loaded: continue - if "learnable_sink_param" in name: - if name in params_dict: - narrow_weight = loaded_weight[head_rank_start:head_rank_end] - n = narrow_weight.shape[0] - with torch.no_grad(): - params_dict[name][:n].copy_(narrow_weight) - loaded_params.add(name) - continue - remapped_name = maybe_remap_kv_scale_name(name, params_dict) if remapped_name is None: continue