iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix][RL] Load attention sinks through weight_loader for layerwise reload (#57797)

Browse Source 
Signed-off-by: Ronald1995 <ronaldautomobile@163.com>
Co-authored-by: aoshen02 <aoshen@inferact.ai>
This commit is contained in:
Ronald
2026-10-14 10:17:58 +00:00
committed by GitHub
co-authored by aoshen02
parent f9b65c30c7
commit 19aba3e4d5
30 changed files with 403 additions and 476 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/model_executor/model_loader/test_reload.py
+204
View File
@@ -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.
tests/model_executor/test_weight_utils.py
+92
View File
@@ -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
tests/models/test_deepseek_v4_dspark_rocm.py
-4
View File
@@ -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
-4
View File
@@ -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",
vllm/model_executor/model_loader/weight_utils.py
+32 -1
View File
@@ -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"""
vllm/model_executor/models/gpt_oss.py
+2 -44
View File
@@ -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,
)
vllm/model_executor/models/laguna.py
+3 -23
View File
@@ -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):
vllm/model_executor/models/mimo_v2.py
+5 -15
View File
@@ -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
vllm/model_executor/models/mimo_v2_mtp.py
-8
View File
@@ -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
)
vllm/model_executor/models/qwen3_dflash.py
+3 -22
View File
@@ -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):
vllm/models/deepseek_v4/amd/dspark.py
-15
View File
@@ -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"
vllm/models/deepseek_v4/amd/model.py
-16
View File
@@ -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
vllm/models/deepseek_v4/amd/mtp.py
-18
View File
@@ -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(
vllm/models/deepseek_v4/attention.py
+2 -4
View File
@@ -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,
vllm/models/deepseek_v4/cpu/model.py
-16
View File
@@ -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
vllm/models/deepseek_v4/nvidia/dspark.py
-15
View File
@@ -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"
vllm/models/deepseek_v4/nvidia/model.py
-17
View File
@@ -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
vllm/models/deepseek_v4/nvidia/mtp.py
-18
View File
@@ -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
vllm/models/deepseek_v4/xpu/dspark.py
-17
View File
@@ -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"
vllm/models/deepseek_v4/xpu/model.py
-16
View File
@@ -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
vllm/models/deepseek_v4/xpu/mtp.py
-18
View File
@@ -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(
vllm/models/deepseek_v41/amd/dspark.py
-15
View File
@@ -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"
vllm/models/deepseek_v41/amd/model.py
-17
View File
@@ -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
vllm/models/deepseek_v41/attention.py
+2 -4
View File
@@ -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,
vllm/models/deepseek_v41/nvidia/dspark.py
-15
View File
@@ -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"
vllm/models/deepseek_v41/nvidia/model.py
-17
View File
@@ -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
vllm/models/hy_v4/amd/model.py
+26 -46
View File
@@ -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
vllm/models/hy_v4/nvidia/attention.py
+6 -7
View File
@@ -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
vllm/models/hy_v4/nvidia/model.py
+26 -46
View File
@@ -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
vllm/models/hy_v4/nvidia/mtp.py
-18
View File
@@ -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.<id>.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