diff --git a/tests/models/kimi_k3/test_weight_loading.py b/tests/models/kimi_k3/test_weight_loading.py deleted file mode 100644 index b24ad17d9a..0000000000 --- a/tests/models/kimi_k3/test_weight_loading.py +++ /dev/null @@ -1,72 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from types import SimpleNamespace - -import pytest -import torch -from torch import nn - -from vllm.models.kimi_k3.nvidia.model import ( - KimiK3ForConditionalGeneration, - KimiLinearForCausalLM, -) - -pytestmark = pytest.mark.cpu_test - - -class _FakeKimiLinearModel(nn.Module): - def __init__(self) -> None: - super().__init__() - self.tensor_a = nn.Parameter(torch.zeros(1)) - self.tensor_c = nn.Parameter(torch.zeros(1)) - self.finalized_values: list[tuple[float, float]] = [] - - def load_weights(self, weights): - params = dict(self.named_parameters()) - loaded = set() - for name, value in weights: - params[name].data.copy_(value) - loaded.add(name) - return loaded - - def finalize_mega_moe_weights(self) -> None: - self.finalized_values.append((self.tensor_a.item(), self.tensor_c.item())) - # MegaMoE finalization replaces its original weight parameters. - self.tensor_a = None - self.tensor_c = None - - -def test_interleaved_composite_weights_finalize_kimi_once_after_loading() -> None: - language_model = object.__new__(KimiLinearForCausalLM) - nn.Module.__init__(language_model) - language_model.config = SimpleNamespace(tie_word_embeddings=False) - language_model.model = _FakeKimiLinearModel() - - model = object.__new__(KimiK3ForConditionalGeneration) - nn.Module.__init__(model) - model.language_model = language_model - model.vision_tower = nn.Module() - model.vision_tower.tensor_b = nn.Parameter(torch.zeros(1)) - - loaded = model.load_weights( - iter( - [ - ("language_model.model.tensor_a", torch.tensor([1.0])), - ("vision_tower.tensor_b", torch.tensor([2.0])), - ("language_model.model.tensor_c", torch.tensor([3.0])), - ] - ) - ) - - assert loaded == { - "language_model.model.tensor_a", - "vision_tower.tensor_b", - "language_model.model.tensor_c", - } - assert language_model.model.finalized_values == [] - - model.process_weights_after_loading() - - assert language_model.model.finalized_values == [(1.0, 3.0)] - assert model.vision_tower.tensor_b.item() == 2.0 diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index 250a9d7291..6f1ff1afab 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -12,9 +12,6 @@ from vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe import ( from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( bind_routed_experts_capturer, ) -from vllm.model_executor.layers.quantization.utils.fp8_utils import ( - deepgemm_post_process_weight_scale_block, -) from vllm.model_executor.models.interfaces import ( get_mixture_of_experts_model, is_mixture_of_experts, @@ -27,7 +24,6 @@ from vllm.models.deepseek_v4.nvidia.model import ( make_deepseek_v4_expert_params_mapping, prepare_mega_gate_routing_metadata, ) -from vllm.models.deepseek_v4.nvidia.mtp import DeepSeekV4MTP from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs from vllm.models.deepseek_v41.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID from vllm.models.deepseek_v41.nvidia.model import DeepseekV4MoE as DeepseekV41MoE @@ -43,6 +39,15 @@ pytestmark = pytest.mark.skipif( ) +@pytest.fixture(autouse=True) +def sm100_mega_moe_backend(monkeypatch): + """Build the SM100 kernel layout on any GPU; kernels are gated per test.""" + monkeypatch.setattr( + "vllm.models.deepseek_v4.nvidia.model.get_deep_gemm_mega_moe_backend", + lambda *args: DeepGemmSm100MegaMoEBackend(), + ) + + @pytest.fixture def v41_moe_config(dist_init): return SimpleNamespace( @@ -303,18 +308,6 @@ def test_deepseek_v4_mega_moe_expert_mapping(): ] -def test_deepseek_v4_mega_moe_ue8m0_uint8_to_float(): - raw = torch.tensor([0, 126, 127, 128], dtype=torch.uint8) - - decoded = DeepseekV4MegaMoEExperts._ue8m0_uint8_to_float(raw) - - assert torch.equal(decoded.view(torch.int32), raw.to(torch.int32) << 23) - assert decoded[0].item() == 0.0 - assert decoded[1].item() == 0.5 - assert decoded[2].item() == 1.0 - assert decoded[3].item() == 2.0 - - @pytest.mark.parametrize("use_kimi", [False, True]) def test_deep_gemm_mega_moe_capture_precedes_eplb(monkeypatch, use_kimi): experts_cls = DeepseekV4MegaMoEExperts @@ -427,16 +420,141 @@ def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership(): return_success=True, ) - assert torch.equal(experts.w13_weight[0, :128], w1) - assert torch.equal(experts.w13_weight[0, 128:], w3) - assert torch.equal(experts.w2_weight[0], w2) + interleaved = experts.w13_weight[0].view(torch.uint8).view(16, 2, 8, 64) + assert torch.equal(interleaved[:, 0], w1.view(16, 8, 64)) + assert torch.equal(interleaved[:, 1], w3.view(16, 8, 64)) + assert torch.equal(experts.w2_weight[0].view(torch.uint8), w2) assert torch.count_nonzero(experts.w13_weight[1]) == 0 +def _reference_mega_moe_layout(w13, s13, w2, s2): + """CPU copy of DeepGEMM's transform_weights_for_mega_moe on packed scales.""" + + def interleave(t): + g, n, k = t.shape + halves = t.view(g, 2, n // 16, 8, k) + return torch.stack([halves[:, 0], halves[:, 1]], dim=2).reshape(g, n, k) + + def utccp(sf): + g, mn, k = sf.shape + return sf.reshape(g, -1, 4, 32, k).transpose(2, 3).reshape(g, mn, k) + + return ( + interleave(w13), + utccp(interleave(s13.view(torch.int32))), + w2, + utccp(s2.view(torch.int32)), + ) + + +def _make_mega_moe_experts(num_local_experts, num_logical_experts): + return DeepseekV4MegaMoEExperts( + SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=4), + compilation_config=SimpleNamespace(static_forward_context={}), + ), + num_experts=num_local_experts, + num_local_experts=num_local_experts, + experts_start_idx=0, + top_k=1, + hidden_size=256, + intermediate_size=384, + num_logical_experts=num_logical_experts, + ) + + +def _load_mega_moe_checkpoint_shuffled(experts, num_logical_experts, seed): + hidden, inter = experts.hidden_size, experts.intermediate_size + shapes = { + "w13_weight": (num_logical_experts, 2 * inter, hidden // 2), + "w13_weight_scale": (num_logical_experts, 2 * inter, hidden // 32), + "w2_weight": (num_logical_experts, hidden, inter // 2), + "w2_weight_scale": (num_logical_experts, hidden, inter // 32), + } + generator = torch.Generator().manual_seed(seed) + ckpt = { + name: torch.randint(0, 256, shape, dtype=torch.uint8, generator=generator) + for name, shape in shapes.items() + } + shards = [] + for expert_id in range(num_logical_experts): + for name, tensor in ckpt.items(): + if name.startswith("w2_"): + shards.append((name, tensor[expert_id], "w2", expert_id)) + else: + shards.append((name, tensor[expert_id, :inter], "w1", expert_id)) + shards.append((name, tensor[expert_id, inter:], "w3", expert_id)) + for i in torch.randperm(len(shards), generator=generator).tolist(): + name, loaded_weight, shard_id, expert_id = shards[i] + if name.endswith("scale"): + loaded_weight = loaded_weight.view(torch.float8_e8m0fnu) + assert experts.weight_loader( + getattr(experts, name), + loaded_weight, + f"experts.{name}", + shard_id=shard_id, + expert_id=expert_id, + return_success=True, + ) + return ckpt + + +def _expected_mega_moe_layout(ckpt, num_local_experts, num_logical_experts): + physical = torch.arange(num_local_experts) % num_logical_experts + ckpt = {name: tensor[physical] for name, tensor in ckpt.items()} + expected = _reference_mega_moe_layout(*ckpt.values()) + if is_deep_gemm_supported(): + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + w13, s13, w2, s2 = (t.cuda() for t in ckpt.values()) + hidden, inter = w2.shape[1], w2.shape[2] * 2 + scales = [ + deep_gemm.transform_sf_into_required_layout( + (sf.to(torch.int32) << 23).view(torch.float32), + mn, + k, + (1, 32), + num_local_experts, + ) + for sf, mn, k in ((s13, 2 * inter, hidden), (s2, hidden, inter)) + ] + (l1, l1_sf), (l2, l2_sf) = deep_gemm.transform_weights_for_mega_moe( + (w13.view(torch.int8), scales[0]), (w2.view(torch.int8), scales[1]) + ) + for actual, reference in zip((l1, l1_sf, l2, l2_sf), expected): + assert torch.equal(actual.cpu().view(reference.dtype), reference) + return expected + + +def _assert_mega_moe_layout(experts, expected): + actual = (*experts._transformed_l1_weights, *experts._transformed_l2_weights) + for tensor, reference in zip(actual, expected): + assert torch.equal(tensor.view(reference.dtype), reference) + + +def test_deepseek_v4_mega_moe_loads_shards_into_kernel_layout(): + """Shards in any order, including EPLB replicas, land in DeepGEMM's layout.""" + experts = _make_mega_moe_experts(num_local_experts=4, num_logical_experts=3) + ckpt = _load_mega_moe_checkpoint_shuffled(experts, 3, seed=0) + _assert_mega_moe_layout(experts, _expected_mega_moe_layout(ckpt, 4, 3)) + # DeepGEMM's MN-major scale strides. + assert experts.w13_weight_scale.stride() == (768 * 2, 1, 768) + assert experts.w2_weight_scale.stride() == (256 * 3, 1, 256) + + +def test_deepseek_v4_mega_moe_reload_overwrites_kernel_layout_in_place(): + """A second load must update the live kernel tensors without reallocating.""" + experts = _make_mega_moe_experts(num_local_experts=2, num_logical_experts=2) + _load_mega_moe_checkpoint_shuffled(experts, 2, seed=0) + data_ptrs = [param.data_ptr() for param in experts.parameters()] + ckpt = _load_mega_moe_checkpoint_shuffled(experts, 2, seed=1) + assert [param.data_ptr() for param in experts.parameters()] == data_ptrs + _assert_mega_moe_layout(experts, _expected_mega_moe_layout(ckpt, 2, 2)) + + @pytest.mark.parametrize("intermediate_size", [512, 2304]) -def test_deepseek_v4_mega_moe_preserves_checkpoint_dimensions( - monkeypatch, intermediate_size -): +def test_deepseek_v4_mega_moe_preserves_checkpoint_dimensions(intermediate_size): """Keep native widths so V4.1 routed and shared experts can fuse.""" experts = DeepseekV4MegaMoEExperts( SimpleNamespace( @@ -450,216 +568,149 @@ def test_deepseek_v4_mega_moe_preserves_checkpoint_dimensions( hidden_size=128, intermediate_size=intermediate_size, ) - originals = [] - for param in ( - experts.w13_weight, - experts.w13_weight_scale, - experts.w2_weight, - experts.w2_weight_scale, - ): - param.data.random_(1, 128) - originals.append(param.data.clone()) - monkeypatch.setattr( - experts, "_ensure_backend", lambda: DeepGemmSm100MegaMoEBackend() - ) - monkeypatch.setattr( - "vllm.utils.deep_gemm._import_deep_gemm", - lambda: SimpleNamespace( - transform_sf_into_required_layout=lambda sf, *args: sf, - transform_weights_for_mega_moe=lambda l1, l2: (l1, l2), - ), - ) - experts.finalize_weights() assert experts.intermediate_size == intermediate_size + assert experts.w13_weight.shape == (2, 2 * intermediate_size, 64) + assert experts.w13_weight_scale.shape == (2, 2 * intermediate_size, 1) + assert experts.w2_weight.shape == (2, 128, intermediate_size // 2) + assert experts.w2_weight_scale.shape == (2, 128, intermediate_size // 128) transformed = (*experts._transformed_l1_weights, *experts._transformed_l2_weights) - for actual, original in zip(transformed, originals): - expected = ( - original - if actual.dtype == torch.int8 - else experts._ue8m0_uint8_to_float(original) - ) - assert torch.equal(actual.view(expected.dtype), expected) + for actual, param in zip(transformed, experts.parameters()): + assert actual.data_ptr() == param.data_ptr() - transformed_l1 = experts._transformed_l1_weights - experts.finalize_weights() - assert experts._transformed_l1_weights is transformed_l1 + +def test_deepseek_v4_mega_moe_loads_shared_experts_into_kernel_layout(monkeypatch): + """Shared expert shards land directly in the MegaMoE kernel layout.""" + monkeypatch.setattr( + DeepGemmSm100MegaMoEBackend, "supports_shared_experts", lambda *args: True + ) + experts = DeepseekV4MegaMoEExperts( + SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=4), + compilation_config=SimpleNamespace(static_forward_context={}), + ), + num_experts=2, + num_local_experts=2, + experts_start_idx=0, + top_k=1, + hidden_size=256, + intermediate_size=384, + num_shared_experts=1, + ) + assert experts.has_fused_shared_experts + hidden, inter = 256, 384 + generator = torch.Generator().manual_seed(42) + + def rand(*shape): + return torch.randint(0, 256, shape, dtype=torch.uint8, generator=generator) + + w1, w3, w2 = rand(inter, hidden), rand(inter, hidden), rand(hidden, inter) + s1, s3, s2 = ( + rand(inter, hidden // 32), + rand(inter, hidden // 32), + rand(hidden, inter // 32), + ) + fp8, e8m0 = torch.float8_e4m3fn, torch.float8_e8m0fnu + # The down shard arrives as `.w2.` (MTP, DSpark) or `.down_proj.` (mapped). + for name, data in [ + ("shared_experts.w1.weight", w1.view(fp8)), + ("shared_experts.w3.weight", w3.view(fp8)), + ("shared_experts.down_proj.weight", w2.view(fp8)), + ("shared_experts.w1.weight_scale", s1.view(e8m0)), + ("shared_experts.w3.weight_scale", s3.view(e8m0)), + ("shared_experts.w2.weight_scale", s2.view(e8m0)), + ]: + loaded = experts.load_shared_expert_weight(f"ffn.{name}", data) + # Weight tracking checks this name against the registered parameters. + assert loaded in {f"ffn.experts.{n}" for n, _ in experts.named_parameters()} + + def l1_halves(): + return experts.shared_l1_weight.view(torch.uint8).view(-1, 2, 8, hidden) + + # Every tensor matches DeepGEMM's batched transform of the checkpoint. + expected = _reference_mega_moe_layout( + *(t.unsqueeze(0) for t in (torch.cat([w1, w3]), torch.cat([s1, s3]), w2, s2)) + ) + actual = ( + experts.shared_l1_weight, + experts.shared_l1_weight_scale, + experts.shared_l2_weight, + experts.shared_l2_weight_scale, + ) + for tensor, reference in zip(actual, expected): + assert torch.equal(tensor.view(reference.dtype), reference[0]) + # A reload of one shard overwrites in place and leaves the other half alone. + ptr = experts.shared_l1_weight.data_ptr() + w1b = rand(inter, hidden) + experts.load_shared_expert_weight("shared_experts.w1.weight", w1b.view(fp8)) + assert experts.shared_l1_weight.data_ptr() == ptr + assert torch.equal(l1_halves()[:, 0], w1b.view(-1, 8, hidden)) + assert torch.equal(l1_halves()[:, 1], w3.view(-1, 8, hidden)) + with pytest.raises(ValueError): + experts.load_shared_expert_weight("shared_experts.gate_up_proj.weight", w1) + + +def _fake_deep_gemm(with_shared: bool): + def old_buffer(group): ... + + def old_kernel(y): ... + + def new_buffer(group, num_shared_experts=0): ... + + def new_kernel(y, shared_l1_weights=None, shared_l2_weights=None): ... + + return SimpleNamespace( + get_symm_buffer_for_mega_moe=new_buffer if with_shared else old_buffer, + fp8_fp4_mega_moe=new_kernel if with_shared else old_kernel, + get_block_m_for_mega_moe=lambda *args: 0, + ) @pytest.mark.parametrize( - "hidden_size,intermediate_size,block_size,mxfp8,packed", + "weight_dtype,scale_dtype,scale_shape,with_shared,fused", [ - pytest.param(128, 512, 128, False, False, id="aligned-block128"), - pytest.param(128, 512, 128, False, True, id="aligned-block128-packed"), - pytest.param(128, 512, 128, True, False, id="aligned-block128-mxfp8"), - pytest.param(5120, 2304, 32, False, False, id="deepseek-v41-flash"), - pytest.param(128, 2304, 128, False, False, id="native-block128"), - pytest.param(128, 2304, 128, False, True, id="native-block128-packed"), - pytest.param(5120, 2304, 32, True, False, id="deepseek-v41-flash-mxfp8"), + (torch.float8_e4m3fn, torch.float8_e8m0fnu, (2, 2), True, True), + (torch.float8_e4m3fn, torch.uint8, (256, 8), True, True), + (torch.float8_e4m3fn, torch.float8_e8m0fnu, (2, 2), False, False), + (torch.bfloat16, None, None, True, False), + (torch.float8_e4m3fn, torch.float32, (2, 2), True, False), + (torch.float8_e4m3fn, torch.uint8, (256, 16), True, False), ], ) -def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights( - monkeypatch, hidden_size, intermediate_size, block_size, mxfp8, packed +def test_deepseek_v4_mega_moe_decides_shared_fusion_at_construction( + monkeypatch, weight_dtype, scale_dtype, scale_shape, with_shared, fused ): - """Shared fusion preserves checkpoint weights and scales at native widths.""" - if packed and not is_deep_gemm_supported(): - pytest.skip("Packing the shared scales requires DeepGEMM") - - class FakeDeepGemm: - transformed_dims: list[tuple[int, int]] = [] - scale_inputs: list[tuple[int, ...]] = [] - scale_values: list[torch.Tensor] = [] - - @staticmethod - def get_symm_buffer_for_mega_moe(*args, num_shared_experts=0, **kwargs): - return None - - @staticmethod - def get_block_m_for_mega_moe(*args, **kwargs): - return 128 - - @staticmethod - def fp8_fp4_mega_moe( - y, - l1_weights, - l2_weights, - sym_buffer, - shared_l1_weights=None, - shared_l2_weights=None, - **kwargs, - ): - return None - - @classmethod - def transform_sf_into_required_layout(cls, sf, mn, k, *args, **kwargs): - cls.scale_inputs.append(tuple(sf.shape)) - cls.scale_values.append(sf) - return torch.empty((sf.shape[0], mn, k // 128), dtype=torch.int32) - - @classmethod - def transform_weights_for_mega_moe(cls, l1_weights, l2_weights): - cls.transformed_dims.append((l1_weights[0].dim(), l2_weights[0].dim())) - if l1_weights[0].dim() == 2: - return (l1_weights[0].clone(), l1_weights[1]), l2_weights - return l1_weights, l2_weights - - vllm_config = SimpleNamespace( - scheduler_config=SimpleNamespace(max_num_batched_tokens=4), - compilation_config=SimpleNamespace(static_forward_context={}), + """Fuse only shared experts the installed kernel and the 1x32 layout fit.""" + monkeypatch.setattr( + "vllm.utils.deep_gemm._import_deep_gemm", lambda: _fake_deep_gemm(with_shared) ) + + def linear(): + scale = ( + None if scale_dtype is None else torch.empty(scale_shape).to(scale_dtype) + ) + return SimpleNamespace( + weight=torch.empty(256, 256, dtype=weight_dtype), weight_scale_inv=scale + ) + + shared = SimpleNamespace(gate_up_proj=linear(), down_proj=linear()) experts = DeepseekV4MegaMoEExperts( - vllm_config, + SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=4), + compilation_config=SimpleNamespace(static_forward_context={}), + ), num_experts=2, - num_local_experts=1, + num_local_experts=2, experts_start_idx=0, top_k=1, - hidden_size=hidden_size, - intermediate_size=intermediate_size, + hidden_size=256, + intermediate_size=128, num_shared_experts=1, - ) - monkeypatch.setattr( - experts, "_ensure_backend", lambda: DeepGemmSm100MegaMoEBackend() + shared_experts=shared, ) - def fp8_parameter(*shape): - return torch.nn.Parameter( - torch.empty(*shape, dtype=torch.uint8) - .random_(1, 119) - .view(torch.float8_e4m3fn), - requires_grad=False, - ) - - def scale_parameter(*shape, dtype=torch.int32): - return torch.nn.Parameter(torch.ones(*shape, dtype=dtype), requires_grad=False) - - shared_experts = SimpleNamespace( - gate_up_proj=SimpleNamespace( - weight=fp8_parameter(2 * intermediate_size, hidden_size), - weight_block_size=(block_size, block_size), - weight_scale_inv=scale_parameter( - 2 * intermediate_size // block_size, - hidden_size // block_size, - dtype=torch.float8_e8m0fnu, - ), - ), - down_proj=SimpleNamespace( - weight=fp8_parameter(hidden_size, intermediate_size), - weight_block_size=(block_size, block_size), - weight_scale_inv=scale_parameter( - hidden_size // block_size, - intermediate_size // block_size, - dtype=torch.float8_e8m0fnu, - ), - ), - ) - if mxfp8: - for linear in (shared_experts.gate_up_proj, shared_experts.down_proj): - linear.weight_block_size = (1, 32) - linear.weight_scale = torch.nn.Parameter( - linear.weight_scale_inv.view(torch.uint8) - .repeat_interleave(block_size, dim=0) - .repeat_interleave(block_size // 32, dim=1), - requires_grad=False, - ) - del linear.weight_scale_inv - originals = [] - for linear in (shared_experts.gate_up_proj, shared_experts.down_proj): - scale = linear.weight_scale if mxfp8 else linear.weight_scale_inv - scale.data.view(torch.uint8).random_(124, 130) - block_m, block_k = linear.weight_block_size - scale_1x32 = ( - experts._ue8m0_uint8_to_float(scale.view(torch.uint8)) - .repeat_interleave(block_m, dim=0) - .repeat_interleave(block_k // 32, dim=1) - ) - originals.append((linear.weight.view(torch.uint8).clone(), scale_1x32)) - if packed: - linear.weight_scale_inv = torch.nn.Parameter( - deepgemm_post_process_weight_scale_block( - scale.data.cuda().unsqueeze(0), - *linear.weight.shape, - quant_block_shape=linear.weight_block_size, - num_groups=1, - ) - .squeeze(0) - .cpu(), - requires_grad=False, - ) - monkeypatch.setattr("vllm.utils.deep_gemm._import_deep_gemm", lambda: FakeDeepGemm) - - original_gate_up_ptr = shared_experts.gate_up_proj.weight.data_ptr() - experts.finalize_weights(shared_experts) - - assert experts.has_fused_shared_experts - assert experts.intermediate_size == intermediate_size - assert FakeDeepGemm.transformed_dims == [(3, 3), (2, 2)] - assert FakeDeepGemm.scale_inputs[-2:] == [ - (1, 2 * intermediate_size, hidden_size // 32), - (1, hidden_size, intermediate_size // 32), - ] - for transformed, actual_scale, (weight, scale) in zip( - ( - experts._transformed_shared_l1_weights, - experts._transformed_shared_l2_weights, - ), - FakeDeepGemm.scale_values[-2:], - originals, - ): - assert torch.equal(transformed[0].view(torch.uint8), weight) - assert torch.equal(actual_scale[0], scale) - assert shared_experts.gate_up_proj.weight.data_ptr() != original_gate_up_ptr - assert ( - experts._transformed_shared_l1_weights[0].data_ptr() - == shared_experts.gate_up_proj.weight.data_ptr() - ) - assert ( - experts._transformed_shared_l2_weights[0].data_ptr() - == shared_experts.down_proj.weight.data_ptr() - ) - - transformed_l1 = experts._transformed_shared_l1_weights - experts.finalize_weights(shared_experts) - assert experts._transformed_shared_l1_weights is transformed_l1 + assert experts.has_fused_shared_experts == fused + assert hasattr(experts, "shared_l1_weight") == fused @pytest.mark.parametrize("fused", [False, True]) @@ -676,10 +727,7 @@ def test_deepseek_v4_mega_moe_does_not_double_add_fused_shared_expert( class FakeExperts(torch.nn.Module): has_fused_shared_experts = fused - _backend = DeepGemmSm100MegaMoEBackend() - - def _ensure_backend(self): - return self._backend + backend = DeepGemmSm100MegaMoEBackend() def forward(self, hidden_states, *args, **kwargs): return torch.ones_like(hidden_states) @@ -696,7 +744,7 @@ def test_deepseek_v4_mega_moe_does_not_double_add_fused_shared_expert( moe.use_native_mega_moe = True moe.gate = FakeGate() moe.experts = FakeExperts() - moe.shared_experts = FakeSharedExperts() + moe.shared_experts = None if fused else FakeSharedExperts() moe.scoring_func = "sqrtsoftplus" moe.n_activated_experts = 1 moe.renormalize = True @@ -715,7 +763,8 @@ def test_deepseek_v4_mega_moe_does_not_double_add_fused_shared_expert( expected = 1 if fused else 3 assert torch.all(output == expected) - assert moe.shared_experts.calls == (0 if fused else 1) + if not fused: + assert moe.shared_experts.calls == 1 @pytest.mark.skipif( @@ -887,35 +936,20 @@ def test_deepseek_v4_mega_moe_stages_shared_scale_tma_layout(shared_block_m): assert torch.all(fused_shared_x_sf[~populated] == -1) -def test_deepseek_v4_pwal_hook_finalizes_mega_moe_and_mhc_broadcast(): +def test_deepseek_v4_pwal_hook_finalizes_mhc_broadcast(): """The loader invokes the model-level PWAL hook for every load format, - so it must finalize megamoe + mhc broadcast weights to cover dummy - load, which skips load_weights().""" + so it must finalize mhc broadcast weights to cover dummy load, which + skips load_weights().""" calls = [] stub = SimpleNamespace( model=SimpleNamespace( - finalize_mega_moe_weights=lambda: calls.append("mega_moe"), finalize_mhc_broadcast_weights=lambda: calls.append("mhc"), ) ) DeepseekV4ForCausalLM.process_weights_after_loading(stub) - assert calls == ["mega_moe", "mhc"] - - -def test_deepseek_v4_drafter_pwal_hooks_finalize_mega_moe(): - """MTP/DSpark drafters load as their own top-level models, so each needs - its own PWAL hook now that the megamoe forward no longer finalizes - weights lazily on first use.""" - calls = [] - mtp = SimpleNamespace(finalize_mega_moe_weights=lambda: calls.append("mtp")) - DeepSeekV4MTP.process_weights_after_loading(mtp) - - dspark = SimpleNamespace(_finalize_moe=lambda: calls.append("dspark")) - DSparkDeepseekV4ForCausalLM.process_weights_after_loading(dspark) - - assert calls == ["mtp", "dspark"] + assert calls == ["mhc"] @pytest.fixture diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index 81019421ed..3ea010bbe5 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -242,7 +242,14 @@ def test_v41_dspark_loads_linear_scales( linear_scale_name=scale_name, pad_shared_expert=False, model=SimpleNamespace( - layers=[SimpleNamespace(ffn=SimpleNamespace(use_native_mega_moe=False))], + layers=[ + SimpleNamespace( + ffn=SimpleNamespace( + use_native_mega_moe=False, + experts=SimpleNamespace(has_fused_shared_experts=False), + ) + ) + ], confidence_head=None, ), named_parameters=lambda: [(runtime_name, param)], diff --git a/vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py b/vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py index 633ced1082..b86820f672 100644 --- a/vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py +++ b/vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py @@ -18,19 +18,15 @@ from vllm.platforms import current_platform logger = init_logger(__name__) -def ue8m0_uint8_to_float(sf: torch.Tensor) -> torch.Tensor: - """Reinterpret a uint8 UE8M0 exponent tensor as float32 scale values.""" - return (sf.to(torch.int32) << 23).view(torch.float32) - - class DeepGemmMegaMoEBackend: """Abstract MegaMoE backend that hides deep_gemm-specific details. A backend advertises its expected hidden-state quantization layout via - ``hidden_quant`` and its expected MMA type via ``mma_type``. It is - responsible for the routed-expert weight transform run at load time, the - optional shared-expert weight transform, and the actual MegaMoE kernel - invocation at inference time. Runtime capability checks live in + ``hidden_quant`` and its expected MMA type via ``mma_type``. It owns the + kernel weight layout: :meth:`create_weights` registers parameters in that + layout and :meth:`load_weight` writes each checkpoint tensor into it, so + loading (and reloading) needs no post-load transform. It also runs the + MegaMoE kernel. Runtime capability checks live in :func:`get_deep_gemm_mega_moe_backend`: obtaining a backend already implies that the current device and shapes are supported. """ @@ -39,49 +35,39 @@ class DeepGemmMegaMoEBackend: mma_type: str = "fp8xfp4" hidden_quant: QuantKey = kMxfp8Dynamic - def transform_weights( + def create_weights( self, + layer: nn.Module, *, - w13_weight: torch.Tensor, - w13_weight_scale: torch.Tensor, - w2_weight: torch.Tensor, - w2_weight_scale: torch.Tensor, num_local_experts: int, hidden_size: int, intermediate_size: int, - activation: str | None = None, - ) -> tuple[ - tuple[torch.Tensor, torch.Tensor], - tuple[torch.Tensor, torch.Tensor], - ]: + num_shared_experts: int, + ) -> None: + """Register the expert (and fused shared-expert) kernel-layout parameters: + ``w13_weight``, ``w13_weight_scale``, ``w2_weight``, ``w2_weight_scale`` + and, if ``num_shared_experts``, ``shared_l{1,2}_weight[_scale]``. Routed + parameters are indexed by local expert on dim 0, which EPLB moves.""" raise NotImplementedError - def supports_shared_experts(self, deep_gemm) -> bool: - return False - - def transform_shared_expert_weights( + def load_weight( self, - *, - shared_experts: nn.Module, - num_shared_experts: int, - hidden_size: int, - intermediate_size: int, - prefix: str, - ) -> ( - tuple[ - tuple[torch.Tensor, torch.Tensor], - tuple[torch.Tensor, torch.Tensor], - ] - | None - ): - """Return ``(l1, l2)`` transformed shared-expert weights, or ``None``. + dst: torch.Tensor, + src: torch.Tensor, + shard_id: str, + is_scale: bool, + ) -> None: + """Write checkpoint shard ``src`` (``w1``/``w2``/``w3``) into ``dst``, one + expert's slice or a shared-expert parameter, in the kernel layout.""" + raise NotImplementedError - ``shared_experts`` must expose ``gate_up_proj`` and ``down_proj`` - linear layers. A ``None`` return means the backend cannot fuse shared - experts for this module; the caller should fall back to a separate - shared MLP. + def supports_shared_experts(self, shared_experts: nn.Module) -> bool: + """Whether ``shared_experts`` can be fused into the MegaMoE kernel. + + Decided at construction from the installed kernel API and the shared + MLP's parameters, so the fused layout is fixed before any weight loads. """ - return None + return False def run_mega_moe( self, @@ -133,229 +119,153 @@ class DeepGemmSm100MegaMoEBackend(DeepGemmMegaMoEBackend): mma_type = "fp8xfp4" hidden_quant = kMxfp8Dynamic - def transform_weights( + def create_weights( self, + layer: nn.Module, *, - w13_weight: torch.Tensor, - w13_weight_scale: torch.Tensor, - w2_weight: torch.Tensor, - w2_weight_scale: torch.Tensor, num_local_experts: int, hidden_size: int, intermediate_size: int, - activation: str | None = None, - ) -> tuple[ - tuple[torch.Tensor, torch.Tensor], - tuple[torch.Tensor, torch.Tensor], - ]: - from vllm.utils.deep_gemm import _import_deep_gemm - - deep_gemm = _import_deep_gemm() - - w13_scale = deep_gemm.transform_sf_into_required_layout( - ue8m0_uint8_to_float(w13_weight_scale).contiguous(), - 2 * intermediate_size, - hidden_size, - (1, 32), - num_local_experts, - ) - w2_scale = deep_gemm.transform_sf_into_required_layout( - ue8m0_uint8_to_float(w2_weight_scale).contiguous(), - hidden_size, - intermediate_size, - (1, 32), - num_local_experts, - ) - kwargs = {} if activation is None else {"activation": activation} - return deep_gemm.transform_weights_for_mega_moe( - (w13_weight.view(torch.int8).contiguous(), w13_scale), - (w2_weight.view(torch.int8).contiguous(), w2_scale), - **kwargs, - ) - - def supports_shared_experts(self, deep_gemm) -> bool: - """Check the Python API before touching a symmetric-memory group. - - This also gives users of an older precompiled vLLM wheel a safe serial - fallback instead of failing halfway through multi-rank buffer setup. - """ - try: - buffer_params = signature(deep_gemm.get_symm_buffer_for_mega_moe).parameters - kernel_params = signature(deep_gemm.fp8_fp4_mega_moe).parameters - except (TypeError, ValueError): - return False - return ( - hasattr(deep_gemm, "get_block_m_for_mega_moe") - and hasattr(deep_gemm, "transform_weights_for_mega_moe") - and "num_shared_experts" in buffer_params - and "shared_l1_weights" in kernel_params - and "shared_l2_weights" in kernel_params - ) - - def transform_shared_expert_weights( - self, - *, - shared_experts: nn.Module, num_shared_experts: int, - hidden_size: int, - intermediate_size: int, - prefix: str, - ) -> ( - tuple[ - tuple[torch.Tensor, torch.Tensor], - tuple[torch.Tensor, torch.Tensor], - ] - | None - ): - from vllm.utils.deep_gemm import _import_deep_gemm + ) -> None: + def register(name: str, data: torch.Tensor) -> None: + layer.register_parameter(name, nn.Parameter(data, requires_grad=False)) - deep_gemm = _import_deep_gemm() - - gate_up = shared_experts.gate_up_proj - down = shared_experts.down_proj - gate_up_weight = gate_up.weight.data - gate_up_scale = ( - gate_up.weight_scale - if hasattr(gate_up, "weight_scale") - else gate_up.weight_scale_inv - ).data - down_weight = down.weight.data - down_scale = ( - down.weight_scale - if hasattr(down, "weight_scale") - else down.weight_scale_inv - ).data - - # MegaMoE's shared FP8 MMA consumes a 1x32 scale for every weight row, - # while the checkpoint uses coarser block-FP8 scales (usually - # 128x128). Build a dedicated, numerically equivalent scale view. - ue8m0_scale_dtypes = (torch.float8_e8m0fnu, torch.uint8, torch.int32) - if ( - gate_up_scale.dtype in ue8m0_scale_dtypes - and down_scale.dtype in ue8m0_scale_dtypes - ): - gate_up_scale = self._prepare_shared_expert_scale( - deep_gemm, - gate_up, - gate_up_scale, - gate_up_weight.shape[0], - gate_up_weight.shape[1], - prefix, - ) - down_scale = self._prepare_shared_expert_scale( - deep_gemm, - down, - down_scale, - down_weight.shape[0], - down_weight.shape[1], - prefix, - ) - - if gate_up_scale is None or down_scale is None: - return None - - shared_intermediate_size = intermediate_size * num_shared_experts - expected_gate_up_shape = ( - 2 * shared_intermediate_size, - hidden_size, + # FP4 weights two per byte; UE8M0 scales MN-major, four per int32 along K. + e, h, i = num_local_experts, hidden_size, intermediate_size + register("w13_weight", torch.zeros(e, 2 * i, h // 2, dtype=torch.int8)) + register( + "w13_weight_scale", + torch.zeros(e, h // 128, 2 * i, dtype=torch.int32).transpose(1, 2), ) - expected_down_shape = (hidden_size, shared_intermediate_size) - if ( - gate_up_weight.dtype != torch.float8_e4m3fn - or down_weight.dtype != torch.float8_e4m3fn - or gate_up_scale.dtype != torch.int32 - or down_scale.dtype != torch.int32 - or tuple(gate_up_weight.shape) != expected_gate_up_shape - or tuple(down_weight.shape) != expected_down_shape - ): - logger.warning( - "Disabling native MegaMoE shared-expert fusion for %s: expected " - "replicated block-FP8 weights with gate_up=%s, down=%s, and " - "DeepGEMM int32 scales; got gate_up=%s/%s/%s and down=%s/%s/%s.", - prefix, - expected_gate_up_shape, - expected_down_shape, - tuple(gate_up_weight.shape), - gate_up_weight.dtype, - gate_up_scale.dtype, - tuple(down_weight.shape), - down_weight.dtype, - down_scale.dtype, - ) - return None - - transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe( - (gate_up_weight, gate_up_scale), - (down_weight, down_scale), + register("w2_weight", torch.zeros(e, h, i // 2, dtype=torch.int8)) + register( + "w2_weight_scale", + torch.zeros(e, i // 128, h, dtype=torch.int32).transpose(1, 2), ) - # L1 interleaving allocates a full copy. Re-home the loader Parameter on - # that storage so the original 2*intermediate*hidden FP8 tensor can be - # released instead of adding roughly 0.7 GiB per rank on DSV4-Flash. - # The generic linear post-load hook may still repack the serial scales, - # but this shared MLP is never called after native fusion is enabled. - gate_up.weight.data = transformed_l1[0] - transformed_l1 = (gate_up.weight.data, transformed_l1[1]) - return transformed_l1, transformed_l2 + if num_shared_experts == 0: + return + # The shared expert is FP8 with the layout of one routed expert slice. + s, fp8 = i * num_shared_experts, torch.float8_e4m3fn + register("shared_l1_weight", torch.zeros(2 * s, h, dtype=fp8)) + register( + "shared_l1_weight_scale", + torch.zeros(h // 128, 2 * s, dtype=torch.int32).transpose(0, 1), + ) + register("shared_l2_weight", torch.zeros(h, s, dtype=fp8)) + register( + "shared_l2_weight_scale", + torch.zeros(s // 128, h, dtype=torch.int32).transpose(0, 1), + ) + + def load_weight( + self, + dst: torch.Tensor, + src: torch.Tensor, + shard_id: str, + is_scale: bool, + ) -> None: + if shard_id not in ("w1", "w2", "w3"): + raise ValueError(f"Unsupported expert shard id: {shard_id}") + rows = dst.shape[0] // (1 if shard_id == "w2" else 2) + if is_scale: + src = self._pack_scale(src, rows, dst.shape[1] * 4) + else: + src = src.view(dst.dtype) + if tuple(src.shape) != (rows, dst.shape[1]): + raise ValueError( + f"MegaMoE {shard_id} weight shape {tuple(src.shape)} does not " + f"match the kernel shard ({rows}, {dst.shape[1]})" + ) + dst, src = self._kernel_layout_shard(dst, src, shard_id, is_scale) + dst.copy_(src) @staticmethod - def _prepare_shared_expert_scale( - deep_gemm, - linear: nn.Module, - scale: torch.Tensor, - mn: int, - k: int, - prefix: str, - ) -> torch.Tensor | None: - block_size = getattr(linear, "weight_block_size", None) - if block_size is None or len(block_size) != 2: - logger.warning( - "Disabling native MegaMoE shared-expert fusion for %s: " - "shared FP8 weight block size is unavailable.", - prefix, + def _pack_scale(sf: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """Expand block UE8M0 scales to one per 1x32 and pack four along K per int32.""" + sf = sf.view(torch.uint8) + if rows % sf.shape[0] or cols % sf.shape[1]: + raise ValueError( + f"Scale shape {tuple(sf.shape)} does not tile ({rows}, {cols})" ) - return None + sf = sf.repeat_interleave(rows // sf.shape[0], 0) + sf = sf.repeat_interleave(cols // sf.shape[1], 1) + return sf.contiguous().clone().view(torch.int32) - block_m, block_k = block_size - num_k_blocks = (k + block_k - 1) // block_k - # The linear hook's DeepGEMM layout packs four k-block bytes per int32. - packed = scale.dtype == torch.int32 - expected_shape = ( - (mn, (num_k_blocks + 3) // 4) - if packed - else ((mn + block_m - 1) // block_m, num_k_blocks) - ) - if block_k % 32 != 0 or tuple(scale.shape) != expected_shape: - logger.warning( - "Disabling native MegaMoE shared-expert fusion for %s: " - "cannot convert shared scale shape %s with block size %s " - "to MegaMoE's 1x32 layout for weight (%d, %d).", - prefix, - tuple(scale.shape), - tuple(block_size), - mn, - k, + @staticmethod + def _kernel_layout_shard( + dst: torch.Tensor, + src: torch.Tensor, + shard_id: str, + is_scale: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Pair a checkpoint shard with its destination in the MegaMoE layout.""" + k = src.shape[-1] + if shard_id == "w2": + if not is_scale: + return dst, src + # UTCCP: scale row 128*b + 32*j + i -> 128*b + 4*i + j. + return dst.view(-1, 32, 4, k).transpose(1, 2), src.reshape(-1, 4, 32, k) + half = 0 if shard_id == "w1" else 1 + if not is_scale: + # Interleave gate/up by 8 rows: r -> (r // 8) * 16 + 8 * half + r % 8. + return dst.view(-1, 2, 8, k)[:, half], src.reshape(-1, 8, k) + # Interleave + UTCCP: 64b + 16j + 8q + m -> 128b + 64q + 32half + 4m + j. + dst = dst.view(-1, 2, 2, 8, 4, k)[:, :, half].permute(0, 3, 1, 2, 4) + return dst, src.reshape(-1, 4, 2, 8, k) + + def supports_shared_experts(self, shared_experts: nn.Module) -> bool: + for linear in (shared_experts.gate_up_proj, shared_experts.down_proj): + weight = linear.weight + scale = getattr(linear, "weight_scale", None) + if scale is None: + scale = getattr(linear, "weight_scale_inv", None) + if not self._tiles_1x32(weight, scale): + logger.warning_once( + "Disabling native MegaMoE shared-expert fusion: needs FP8 " + "weights with UE8M0 block scales that tile to 1x32; got " + "weight %s/%s and scale %s/%s.", + tuple(weight.shape), + weight.dtype, + None if scale is None else tuple(scale.shape), + None if scale is None else scale.dtype, + ) + return False + + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + try: + api = { + *signature(deep_gemm.get_symm_buffer_for_mega_moe).parameters, + *signature(deep_gemm.fp8_fp4_mega_moe).parameters, + } + except (AttributeError, TypeError, ValueError): + api = set() + if not ( + hasattr(deep_gemm, "get_block_m_for_mega_moe") + and {"num_shared_experts", "shared_l1_weights", "shared_l2_weights"} <= api + ): + logger.warning_once( + "Disabling native MegaMoE shared-expert fusion because the " + "installed DeepGEMM Python API is older than the vLLM " + "source. Rebuild the vendored _deep_gemm_C extension to enable it.", ) - return None + return False + return True - if packed: - row_scale = scale.flatten().view(torch.uint8).view(mn, -1)[:, :num_k_blocks] - else: - row_scale = scale.view(torch.uint8).repeat_interleave(block_m, dim=0)[:mn] - scale_1x32 = ( - ue8m0_uint8_to_float(row_scale) - .repeat_interleave(block_k // 32, dim=1)[:, : k // 32] - .contiguous() + @staticmethod + def _tiles_1x32(weight: torch.Tensor, scale: torch.Tensor | None) -> bool: + if scale is None or weight.dim() != 2 or scale.dim() != 2: + return False + (rows, cols), (scale_rows, scale_cols) = weight.shape, scale.shape + return ( + weight.dtype == torch.float8_e4m3fn + and scale.dtype in (torch.uint8, torch.float8_e8m0fnu) + and rows % scale_rows == 0 + and cols % scale_cols == 0 + and (cols // scale_cols) % 32 == 0 ) - # The grouped API is used with a singleton dimension to request the - # MN-major, TMA-aligned packed-UE8M0 strides, then squeezed back to the - # 2D layout required for a shared expert. - return deep_gemm.transform_sf_into_required_layout( - scale_1x32.unsqueeze(0), - mn, - k, - (1, 32), - 1, - ).squeeze(0) def run_mega_moe( self, diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index 2a40412e93..db47b7cf11 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -46,7 +46,7 @@ from vllm.model_executor.models.qwen3_dspark import ( DSparkConfidenceHead, DSparkMarkovHead, ) -from vllm.model_executor.models.utils import maybe_prefix +from vllm.model_executor.models.utils import extract_layer_index, maybe_prefix from vllm.models.common.ops.sequence_parallel import ( sp_all_gather, sp_padding_mask, @@ -507,6 +507,13 @@ class DSparkDeepseekV4ForCausalLM(nn.Module, DeepseekV4MixtureOfExperts): else ".weight_scale_inv" ) name = name.removesuffix(".scale") + suffix + if ".shared_experts." in name: + experts = self.model.layers[extract_layer_index(name)].ffn.experts + if getattr(experts, "has_fused_shared_experts", False): + loaded_params.add( + experts.load_shared_expert_weight(name, loaded_weight) + ) + continue if ".shared_experts.w2" in name: name = name.replace(".shared_experts.w2", ".shared_experts.down_proj") if self.pad_shared_expert and ".shared_experts." in name: @@ -574,17 +581,9 @@ class DSparkDeepseekV4ForCausalLM(nn.Module, DeepseekV4MixtureOfExperts): if self.model.confidence_head is not None and not loaded_confidence_head: self.model.confidence_head = None - self.process_weights_after_loading() logger.info_once("DSpark draft model loaded: %d params", len(loaded_params)) return loaded_params - def _finalize_moe(self) -> None: - for layer in self.model.layers: - layer.ffn.finalize_mega_moe_weights() - - def process_weights_after_loading(self) -> None: - self._finalize_moe() - def _remap_dspark_name(self, name: str) -> str | None: """Map a checkpoint ``mtp.{i}.*`` name to this model's parameter path. diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index ff4470dc4c..6dd49f9409 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -33,9 +33,7 @@ from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe import ( - DeepGemmMegaMoEBackend, get_deep_gemm_mega_moe_backend, - ue8m0_uint8_to_float, ) from vllm.model_executor.layers.fused_moe.router.base_router import ( eplb_map_to_physical_and_record, @@ -225,6 +223,7 @@ class DeepseekV4MegaMoEExperts(nn.Module): hidden_size: int, intermediate_size: int, num_shared_experts: int = 0, + shared_experts: nn.Module | None = None, prefix: str = "", num_logical_experts: int | None = None, ): @@ -238,7 +237,6 @@ class DeepseekV4MegaMoEExperts(nn.Module): self.top_k = top_k self.hidden_size = hidden_size self.intermediate_size = intermediate_size - self.num_shared_experts = num_shared_experts self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens self.num_logical_experts = ( @@ -247,64 +245,33 @@ class DeepseekV4MegaMoEExperts(nn.Module): self.eplb_state = EplbLayerState() - weight_attrs = {"weight_loader": self.weight_loader} - self.w13_weight = nn.Parameter( - torch.zeros( - num_local_experts, - 2 * intermediate_size, - hidden_size // 2, - dtype=torch.uint8, - ), - requires_grad=False, + device = torch.get_default_device() + if device.type != current_platform.device_type: # meta under ipc_cache + device = torch.device( + current_platform.device_type, torch.accelerator.current_device_index() + ) + self.backend = get_deep_gemm_mega_moe_backend( + device, hidden_size, intermediate_size ) - set_weight_attrs(self.w13_weight, weight_attrs) - - self.w13_weight_scale = nn.Parameter( - torch.zeros( - num_local_experts, - 2 * intermediate_size, - hidden_size // 32, - dtype=torch.uint8, - ), - requires_grad=False, + if num_shared_experts and not self.backend.supports_shared_experts( + shared_experts + ): + num_shared_experts = 0 + self.num_shared_experts = num_shared_experts + self.backend.create_weights( + self, + num_local_experts=num_local_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_shared_experts=num_shared_experts, ) - set_weight_attrs(self.w13_weight_scale, weight_attrs) + for name, param in self.named_parameters(recurse=False): + shared = name.startswith("shared_") + loader = self.shared_weight_loader if shared else self.weight_loader + set_weight_attrs(param, {"weight_loader": loader}) self.w13_weight_scale.quant_method = "block" - - self.w2_weight = nn.Parameter( - torch.zeros( - num_local_experts, - hidden_size, - intermediate_size // 2, - dtype=torch.uint8, - ), - requires_grad=False, - ) - set_weight_attrs(self.w2_weight, weight_attrs) - - self.w2_weight_scale = nn.Parameter( - torch.zeros( - num_local_experts, - hidden_size, - intermediate_size // 32, - dtype=torch.uint8, - ), - requires_grad=False, - ) - set_weight_attrs(self.w2_weight_scale, weight_attrs) self.w2_weight_scale.quant_method = "block" - self._transformed_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None - self._transformed_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None - self._transformed_shared_l1_weights: ( - tuple[torch.Tensor, torch.Tensor] | None - ) = None - self._transformed_shared_l2_weights: ( - tuple[torch.Tensor, torch.Tensor] | None - ) = None - - self._backend: DeepGemmMegaMoEBackend | None = None - # Register in the static forward context so the custom-op wrapper # can look up this module by name from within a torch.compile graph. compilation_config = vllm_config.compilation_config @@ -332,144 +299,57 @@ class DeepseekV4MegaMoEExperts(nn.Module): return_success: bool = False, ) -> bool | None: local_expert_ids = self._map_global_expert_id(expert_id) - if not local_expert_ids: + param_prefix = "w2_" if shard_id == "w2" else "w13_" + if not local_expert_ids or param_prefix not in weight_name: return False if return_success else None - - loaded_any = False + is_scale = weight_name.endswith("weight_scale") for local_expert_id in local_expert_ids: - expert_data = param.data[local_expert_id] - if shard_id in ("w1", "w3"): - if "w13_" not in weight_name: - continue - shard_offset = 0 if shard_id == "w1" else self.intermediate_size - expert_data = expert_data.narrow( - 0, shard_offset, self.intermediate_size - ) - elif shard_id == "w2": - if "w2_" not in weight_name: - continue - else: - raise ValueError(f"Unsupported expert shard id: {shard_id}") - - if expert_data.shape != loaded_weight.shape: - raise ValueError( - f"DeepSeek V4 MegaMoE expert weight shape mismatch for " - f"{weight_name}: parameter shard {tuple(expert_data.shape)} " - f"vs checkpoint {tuple(loaded_weight.shape)}" - ) - expert_data.copy_(loaded_weight) - loaded_any = True - - if return_success: - return loaded_any - return None - - @staticmethod - def _ue8m0_uint8_to_float(sf: torch.Tensor) -> torch.Tensor: - return ue8m0_uint8_to_float(sf) - - def _ensure_backend(self) -> DeepGemmMegaMoEBackend: - """Lazily resolve the MegaMoE backend. - - The backend is looked up via :func:`get_deep_gemm_mega_moe_backend`, which - performs all runtime capability checks. Failure propagates directly - to the caller. This must be called before :meth:`finalize_weights` - releases the loader-side parameters. - """ - if self._backend is not None: - return self._backend - assert self.w13_weight is not None, ( - "_ensure_backend() must be called before finalize_weights() " - "releases the loader-side parameters." - ) - self._backend = get_deep_gemm_mega_moe_backend( - self.w13_weight.device, - self.hidden_size, - self.intermediate_size, - ) - return self._backend - - def _check_runtime_supported(self) -> None: - """Preserved for API compatibility; delegates to the backend check.""" - self._ensure_backend() - - def _finalize_shared_expert_weights(self, shared_experts: DeepseekV4MLP) -> None: - """Delegate shared-expert weight transform to the backend and update state. - - On success, populates ``self._transformed_shared_l{1,2}_weights``. If - the backend cannot fuse shared experts (returns ``None``), disables - native fusion by setting ``self.num_shared_experts = 0``. - """ - backend = self._ensure_backend() - result = backend.transform_shared_expert_weights( - shared_experts=shared_experts, - num_shared_experts=self.num_shared_experts, - hidden_size=self.hidden_size, - intermediate_size=self.intermediate_size, - prefix=self.prefix, - ) - if result is None: - self.num_shared_experts = 0 - return - ( - self._transformed_shared_l1_weights, - self._transformed_shared_l2_weights, - ) = result - - def finalize_weights(self, shared_experts: DeepseekV4MLP | None = None) -> None: - backend = self._ensure_backend() - from vllm.utils.deep_gemm import _import_deep_gemm - - deep_gemm = _import_deep_gemm() - if self._transformed_l1_weights is None: - self._transformed_l1_weights, self._transformed_l2_weights = ( - backend.transform_weights( - w13_weight=self.w13_weight.data, - w13_weight_scale=self.w13_weight_scale.data, - w2_weight=self.w2_weight.data, - w2_weight_scale=self.w2_weight_scale.data, - num_local_experts=self.num_local_experts, - hidden_size=self.hidden_size, - intermediate_size=self.intermediate_size, - ) + self.backend.load_weight( + param.data[local_expert_id], loaded_weight, shard_id, is_scale ) - # Drop the original loader-side parameters: the MegaMoE kernels only - # consume the transformed views above. transform_weights_for_mega_moe - # allocates a fresh tensor for the L1 weight (see - # _interleave_l1_weights) and fresh SF tensors for L1/L2; the L2 - # weight is the only tensor that aliases the original storage, and - # _transformed_l2_weights still holds it, so the storage stays live - # after we drop the Parameter. - self.w13_weight = None - self.w13_weight_scale = None - self.w2_weight = None - self.w2_weight_scale = None + return True if return_success else None - if shared_experts is None or self.num_shared_experts == 0: - return - if self._transformed_shared_l1_weights is not None: - return - if not backend.supports_shared_experts(deep_gemm): - logger.warning_once( - "Disabling native MegaMoE shared-expert fusion because the " - "installed DeepGEMM Python API is older than the vLLM " - "source. Rebuild the vendored _deep_gemm_C extension to enable it.", - ) - self.num_shared_experts = 0 - return + def load_shared_expert_weight(self, name: str, loaded_weight: torch.Tensor) -> str: + """Load one shared-expert checkpoint tensor; return the parameter name.""" + if ".w2." in name or ".down_proj." in name: + param_name = "shared_l2_weight" + elif ".w1." in name or ".w3." in name: + param_name = "shared_l1_weight" + else: + raise ValueError(f"Unexpected shared-expert tensor: {name}") + if "weight_scale" in name: + param_name += "_scale" + param = getattr(self, param_name) + param.weight_loader(param, loaded_weight, name) + return f"{name.partition('shared_experts.')[0]}experts.{param_name}" - self._finalize_shared_expert_weights(shared_experts) + def shared_weight_loader( + self, param: nn.Parameter, loaded_weight: torch.Tensor, weight_name: str + ) -> None: + shard_id = ( + "w1" if ".w1." in weight_name else "w3" if ".w3." in weight_name else "w2" + ) + is_scale = "weight_scale" in weight_name + self.backend.load_weight(param.data, loaded_weight, shard_id, is_scale) + + @property + def _transformed_l1_weights(self) -> tuple[torch.Tensor, torch.Tensor]: + return self.w13_weight, self.w13_weight_scale + + @property + def _transformed_l2_weights(self) -> tuple[torch.Tensor, torch.Tensor]: + return self.w2_weight, self.w2_weight_scale @property def has_fused_shared_experts(self) -> bool: - return self._transformed_shared_l1_weights is not None + return self.num_shared_experts > 0 def get_symm_buffer(self): from vllm.utils.deep_gemm import _import_deep_gemm deep_gemm = _import_deep_gemm() - backend = self._ensure_backend() + backend = self.backend group = get_ep_group().device_group device = torch.accelerator.current_device_index() key = ( @@ -515,10 +395,6 @@ class DeepseekV4MegaMoEExperts(nn.Module): ) def get_expert_weights(self) -> list[torch.Tensor]: - self.finalize_weights() - assert self._transformed_l1_weights is not None - assert self._transformed_l2_weights is not None - def _to_eplb_view(name: str, t: torch.Tensor) -> torch.Tensor: """Return a (num_local_experts, -1) view with contiguous memory layout.""" assert t.shape[0] == self.num_local_experts @@ -601,7 +477,7 @@ class DeepseekV4MegaMoEExperts(nn.Module): else None, ) - backend = self._ensure_backend() + backend = self.backend symm_buffer = self.get_symm_buffer() shared_x_sf = None shared_block_m = None @@ -630,28 +506,21 @@ class DeepseekV4MegaMoEExperts(nn.Module): hidden_quant=backend.hidden_quant, ) - assert self._transformed_l1_weights is not None - assert self._transformed_l2_weights is not None - if self.has_fused_shared_experts: - backend.run_mega_moe( - y=y, - l1_weights=self._transformed_l1_weights, - l2_weights=self._transformed_l2_weights, - symm_buffer=symm_buffer, - shared_l1_weights=self._transformed_shared_l1_weights, - shared_l2_weights=self._transformed_shared_l2_weights, - activation_clamp=activation_clamp, - fast_math=fast_math, - ) - else: - backend.run_mega_moe( - y=y, - l1_weights=self._transformed_l1_weights, - l2_weights=self._transformed_l2_weights, - symm_buffer=symm_buffer, - activation_clamp=activation_clamp, - fast_math=fast_math, - ) + fused = self.has_fused_shared_experts + backend.run_mega_moe( + y=y, + l1_weights=self._transformed_l1_weights, + l2_weights=self._transformed_l2_weights, + symm_buffer=symm_buffer, + shared_l1_weights=( + (self.shared_l1_weight, self.shared_l1_weight_scale) if fused else None + ), + shared_l2_weights=( + (self.shared_l2_weight, self.shared_l2_weight_scale) if fused else None + ), + activation_clamp=activation_clamp, + fast_math=fast_math, + ) return y @@ -842,9 +711,12 @@ class DeepseekV4MoE(nn.Module): hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, num_shared_experts=(self.n_shared_experts if fuse_shared_experts else 0), + shared_experts=self.shared_experts, prefix=f"{prefix}.experts", ) self.experts = DeepseekV4MegaMoEExperts(vllm_config, **expert_kwargs) + if self.experts.has_fused_shared_experts: + self.shared_experts = None def _init_fused_moe_experts( self, @@ -911,12 +783,7 @@ class DeepseekV4MoE(nn.Module): return self._forward_fused_moe(hidden_states, input_ids) org_shape = hidden_states.shape - # The MegaMoE experts module owns the backend; resolve and assert it. - self.experts._ensure_backend() - assert self.experts._backend is not None, ( - "MegaMoE backend must be resolved before routing." - ) - mega_backend = self.experts._backend + mega_backend = self.experts.backend # Small local padded batches favor GateLinear; 128-expert gates cross earlier. gate_threshold = 1 if self.gate.weight.shape[0] == 128 else 16 if ( @@ -983,10 +850,7 @@ class DeepseekV4MoE(nn.Module): activation_clamp=activation_clamp, ) - if ( - self.shared_experts is not None - and not self.experts.has_fused_shared_experts - ): + if self.shared_experts is not None: shared_output = self.shared_experts(hidden_states) final_hidden_states += shared_output @@ -1004,10 +868,6 @@ class DeepseekV4MoE(nn.Module): return final_hidden_states.view(org_shape) - def finalize_mega_moe_weights(self) -> None: - if self.use_native_mega_moe: - self.experts.finalize_weights(self.shared_experts) - def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]: """Pick the CUDA sparse-MLA attention class for the configured backend. @@ -1559,6 +1419,15 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): ) for name, loaded_weight in weights: + if ".shared_experts." in name: + if is_pp_missing_parameter(name, self): + continue + experts = self.layers[extract_layer_index(name)].ffn.experts + if getattr(experts, "has_fused_shared_experts", False): + loaded_params.add( + experts.load_shared_expert_weight(name, loaded_weight) + ) + continue if pad_shared_expert and ".shared_experts." in name: loaded_weight = self._pad_shared_expert_weight( self.quant_config, name, loaded_weight @@ -1688,10 +1557,6 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): num_experts=self.config.n_routed_experts, ) - def finalize_mega_moe_weights(self) -> None: - for layer in islice(self.layers, self.start_layer, self.end_layer): - layer.ffn.finalize_mega_moe_weights() - def finalize_mhc_broadcast_weights(self) -> None: if not get_pp_group().is_first_rank or self.start_layer >= self.end_layer: return @@ -1891,7 +1756,6 @@ class DeepseekV4ForCausalLM( return loaded_params def process_weights_after_loading(self) -> None: - self.model.finalize_mega_moe_weights() self.model.finalize_mhc_broadcast_weights() def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: diff --git a/vllm/models/deepseek_v4/nvidia/mtp.py b/vllm/models/deepseek_v4/nvidia/mtp.py index c0f3471dd5..0956ab053f 100644 --- a/vllm/models/deepseek_v4/nvidia/mtp.py +++ b/vllm/models/deepseek_v4/nvidia/mtp.py @@ -436,6 +436,13 @@ class DeepSeekV4MTP(nn.Module): else ".weight_scale_inv" ) name = name.removesuffix(".scale") + suffix + if ".shared_experts." in name: + experts = self.model.layers[str(spec_layer)].mtp_block.ffn.experts + if getattr(experts, "has_fused_shared_experts", False): + loaded_params.add( + experts.load_shared_expert_weight(name, loaded_weight) + ) + continue if ".shared_experts.w2" in name: name = name.replace(".shared_experts.w2", ".shared_experts.down_proj") if self.pad_shared_expert and ".shared_experts." in name: @@ -529,17 +536,9 @@ class DeepSeekV4MTP(nn.Module): f"Use a checkpoint that includes MTP layer weights, " f"or disable speculative decoding." ) - self.process_weights_after_loading() logger.info_once("MTP draft model loaded: %d params", len(loaded_params)) return loaded_params - def finalize_mega_moe_weights(self) -> None: - for layer in self.model.layers.values(): - layer.mtp_block.ffn.finalize_mega_moe_weights() - - def process_weights_after_loading(self) -> None: - self.finalize_mega_moe_weights() - def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: """Rewrite the weight name to match the format of the original model. Add .mtp_block for modules in transformer layer block for spec layer diff --git a/vllm/models/deepseek_v41/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py index a09c3159c2..0e9c43d64a 100644 --- a/vllm/models/deepseek_v41/nvidia/dspark.py +++ b/vllm/models/deepseek_v41/nvidia/dspark.py @@ -44,7 +44,7 @@ from vllm.model_executor.models.qwen3_dspark import ( DSparkConfidenceHead, DSparkMarkovHead, ) -from vllm.model_executor.models.utils import maybe_prefix +from vllm.model_executor.models.utils import extract_layer_index, maybe_prefix from vllm.models.common.ops.sequence_parallel import ( sp_all_gather, sp_padding_mask, @@ -431,6 +431,13 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): else f".{self.linear_scale_name}" ) name = name.removesuffix(".scale") + suffix + if ".shared_experts." in name: + experts = self.model.layers[extract_layer_index(name)].ffn.experts + if getattr(experts, "has_fused_shared_experts", False): + loaded_params.add( + experts.load_shared_expert_weight(name, loaded_weight) + ) + continue if ".shared_experts.w2" in name: name = name.replace(".shared_experts.w2", ".shared_experts.down_proj") if self.pad_shared_expert and ".shared_experts." in name: @@ -492,17 +499,9 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): if self.model.confidence_head is not None and not loaded_confidence_head: self.model.confidence_head = None - self.process_weights_after_loading() logger.info_once("DSpark draft model loaded: %d params", len(loaded_params)) return loaded_params - def _finalize_moe(self) -> None: - for layer in self.model.layers: - layer.ffn.finalize_mega_moe_weights() - - def process_weights_after_loading(self) -> None: - self._finalize_moe() - def _remap_dspark_name(self, name: str) -> str | None: """Map a checkpoint ``mtp.{i}.*`` name to this model's parameter path. diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index 8ec0155f73..bc516c8de7 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -1275,6 +1275,15 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): # Vision weights are loaded by the outer multimodal wrapper. logger.warning_once("Skipping non-text weight: %s", name) continue + if ".shared_experts." in name: + if is_pp_missing_parameter(name, self): + continue + experts = self.layers[extract_layer_index(name)].ffn.experts + if getattr(experts, "has_fused_shared_experts", False): + loaded_params.add( + experts.load_shared_expert_weight(name, loaded_weight) + ) + continue if pad_shared_expert and ".shared_experts." in name: loaded_weight = self._pad_shared_expert_weight( self.quant_config, name, loaded_weight @@ -1406,10 +1415,6 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): num_experts=self.config.n_routed_experts, ) - def finalize_mega_moe_weights(self) -> None: - for layer in islice(self.layers, self.start_layer, self.end_layer): - layer.ffn.finalize_mega_moe_weights() - def finalize_mhc_broadcast_weights(self) -> None: if not get_pp_group().is_first_rank or self.start_layer >= self.end_layer: return @@ -1680,7 +1685,6 @@ class DeepseekV41LLMForCausalLM( return loaded_params def process_weights_after_loading(self) -> None: - self.model.finalize_mega_moe_weights() self.model.finalize_mhc_broadcast_weights() def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index ae9b7da3cd..6434ab9b48 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -95,10 +95,7 @@ from vllm.models.common.ops.sequence_parallel import ( sp_reduce_scatter, sp_shard, ) -from vllm.models.deepseek_v4.nvidia.model import ( - DeepseekV4MegaMoEExperts, - DeepseekV4MLP, -) +from vllm.models.deepseek_v4.nvidia.model import DeepseekV4MegaMoEExperts from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention from vllm.models.kimi_k3.nvidia.latent_moe_runner import ( @@ -351,10 +348,6 @@ class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts): self.activation = activation self.activation_beta = activation_beta self.activation_linear_beta = activation_linear_beta - self.register_buffer("_mega_l1_packed", None, persistent=False) - self.register_buffer("_mega_l1_scale", None, persistent=False) - self.register_buffer("_mega_l2_packed", None, persistent=False) - self.register_buffer("_mega_l2_scale", None, persistent=False) def synchronize_first_launch(self) -> None: ep_group = get_ep_group() @@ -366,51 +359,11 @@ class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts): torch.distributed.barrier(group=ep_group.cpu_group) self._synchronized_ep_groups.add(key) - def finalize_weights(self, shared_experts: DeepseekV4MLP | None = None) -> None: - if self._transformed_l1_weights is not None: - return - - # Weight cache IPC engine: the daemon exported the transformed - # buffers; reuse them zero-copy and drop the raw packed params. - if self._mega_l1_packed is not None: - self._transformed_l1_weights = (self._mega_l1_packed, self._mega_l1_scale) - self._transformed_l2_weights = (self._mega_l2_packed, self._mega_l2_scale) - self._drop_raw_mega_weights() - return - - self._check_runtime_supported() - backend = self._ensure_backend() - self._transformed_l1_weights, self._transformed_l2_weights = ( - backend.transform_weights( - w13_weight=self.w13_weight.data, - w13_weight_scale=self.w13_weight_scale.data, - w2_weight=self.w2_weight.data, - w2_weight_scale=self.w2_weight_scale.data, - num_local_experts=self.num_local_experts, - hidden_size=self.hidden_size, - intermediate_size=self.intermediate_size, - activation=self.activation, - ) - ) - l1_packed, l1_scale = self._transformed_l1_weights - l2_packed, l2_scale = self._transformed_l2_weights - self.register_buffer("_mega_l1_packed", l1_packed, persistent=False) - self.register_buffer("_mega_l1_scale", l1_scale, persistent=False) - self.register_buffer("_mega_l2_packed", l2_packed, persistent=False) - self.register_buffer("_mega_l2_scale", l2_scale, persistent=False) - self._drop_raw_mega_weights() - - def _drop_raw_mega_weights(self) -> None: - self.w13_weight = None - self.w13_weight_scale = None - self.w2_weight = None - self.w2_weight_scale = None - def get_symm_buffer(self): from vllm.utils.deep_gemm import _import_deep_gemm deep_gemm = _import_deep_gemm() - backend = self._ensure_backend() + backend = self.backend group = get_ep_group().device_group device = torch.accelerator.current_device_index() key = ( @@ -485,7 +438,7 @@ class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts): else None, ) - backend = self._ensure_backend() + backend = self.backend symm_buffer = self.get_symm_buffer() prepare_megamoe_inputs( hidden_states, @@ -498,9 +451,6 @@ class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts): is_padding=is_padding, hidden_quant=backend.hidden_quant, ) - self.finalize_weights() - assert self._transformed_l1_weights is not None - assert self._transformed_l2_weights is not None backend.run_mega_moe( y=y, l1_weights=self._transformed_l1_weights, @@ -1573,11 +1523,6 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant): loaded_params.add(name) return loaded_params - def finalize_mega_moe_weights(self) -> None: - for module in self.modules(): - if isinstance(module, KimiMoE) and module.use_mega_moe: - module.experts.finalize_weights() - class KimiLinearForCausalLM( nn.Module, @@ -1700,14 +1645,6 @@ class KimiLinearForCausalLM( loader = AutoWeightsLoader(self) return loader.load_weights(weights) - def process_weights_after_loading(self) -> None: - # A parent AutoWeightsLoader may invoke load_weights repeatedly for - # non-contiguous streamed prefixes. Finalize only after the full stream. - self.model.finalize_mega_moe_weights() - # The fused MultiHeadLatentAttention's process_weights_after_loading - # (W_UK_T / W_UV absorption) is driven by the loader's generic post-load - # hook for any AttentionLayerBase, so no manual trigger is needed here. - def get_spec_layer_idx_from_weight_name( config: KimiLinearConfig, weight_name: str @@ -2154,6 +2091,3 @@ class KimiK3ForConditionalGeneration( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) - - def process_weights_after_loading(self) -> None: - self.language_model.process_weights_after_loading() diff --git a/vllm/models/kimi_k3/nvidia/mtp.py b/vllm/models/kimi_k3/nvidia/mtp.py index 038c8db6b3..06a0d33778 100644 --- a/vllm/models/kimi_k3/nvidia/mtp.py +++ b/vllm/models/kimi_k3/nvidia/mtp.py @@ -421,11 +421,6 @@ class KimiK3MTP(nn.Module): f"MTP layer weights, or disable speculative decoding." ) - if use_mega_moe: - for module in self.modules(): - if isinstance(module, KimiMoE) and module.use_mega_moe: - module.experts.finalize_weights() - return loaded_params def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: