[Model][MoE] Load MegaMoE experts directly into the DeepGEMM layout (#60096)
Signed-off-by: aoshen02 <aoshen@inferact.ai> Signed-off-by: Aaron Hao <hao-aaron@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Aaron Hao <hao-aaron@users.noreply.github.com>
This commit is contained in:
aoshen02
co-authored by
Claude Opus 5.5
Aaron Hao
Claude Opus 5.5
Aaron Hao
parent
624a5ddafd
commit
c5bda658d2
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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)],
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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]]:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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]]:
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user