From 4d58fa254a869a5d9d0aab250ab125503d3cef26 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Thu, 15 Oct 2026 12:02:01 -0700 Subject: [PATCH] [MoE] Unify DeepGEMM MegaMoE experts behind one feature-matrix layer; add NVFP4 (#61538) Signed-off-by: Yongye Zhu Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: Kevin H. Luu --- .buildkite/test_areas/kernels.yaml | 4 + tests/kernels/moe/test_deep_gemm_mega_moe.py | 1274 +++++++++++++++++ tests/model_executor/test_deep_gemm_warmup.py | 9 +- tests/models/test_deepseek_v4_mega_moe.py | 738 +--------- .../layers/fused_moe/deep_gemm_mega_moe.py | 1013 +++++++++---- .../layers/fused_moe/prepare_megamoe.py | 682 +++++++++ .../model_executor/warmup/deep_gemm_warmup.py | 11 +- vllm/models/deepseek_v4/nvidia/model.py | 371 +---- .../deepseek_v4/nvidia/ops/prepare_megamoe.py | 323 ----- vllm/models/kimi_k3/nvidia/model.py | 164 +-- 10 files changed, 2753 insertions(+), 1836 deletions(-) create mode 100644 tests/kernels/moe/test_deep_gemm_mega_moe.py create mode 100644 vllm/model_executor/layers/fused_moe/prepare_megamoe.py delete mode 100644 vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 126fff9bdc..901d3251c2 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -138,11 +138,15 @@ steps: - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu - vllm/models/deepseek_v4/common/ops/ - vllm/models/deepseek_v4/nvidia/ + - vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py + - vllm/model_executor/layers/fused_moe/prepare_megamoe.py - vllm/utils/deep_gemm.py - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py + - tests/kernels/moe/test_deep_gemm_mega_moe.py - tests/models/test_deepseek_v4_mega_moe.py commands: - pytest -v -s kernels/test_fused_deepseek_v4_*.py + - pytest -v -s kernels/moe/test_deep_gemm_mega_moe.py - pytest -v -s models/test_deepseek_v4_mega_moe.py mirror: amd: diff --git a/tests/kernels/moe/test_deep_gemm_mega_moe.py b/tests/kernels/moe/test_deep_gemm_mega_moe.py new file mode 100644 index 0000000000..9382ec1c08 --- /dev/null +++ b/tests/kernels/moe/test_deep_gemm_mega_moe.py @@ -0,0 +1,1274 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepGEMM MegaMoE experts: kernel-layout weight loading, shared-expert fusion, +input staging, and numerics against a dequantized reference for every routed +format, activation and shared-expert fusion in the feature matrix.""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe import ( + MEGA_MOE_BF16, + MEGA_MOE_MXFP4, + MEGA_MOE_NVFP4, + MEGA_MOE_WEIGHT_FORMATS, + DeepGemmMegaMoEExperts, + MegaMoEActivation, + MegaMoEWeightSpec, +) +from vllm.model_executor.layers.fused_moe.prepare_megamoe import ( + prepare_megamoe_inputs, + prepare_megamoe_nvfp4_inputs, +) +from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( + bind_routed_experts_capturer, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp8Dynamic, + kNvfp4DynamicToken, +) +from vllm.platforms import current_platform +from vllm.utils.deep_gemm import is_deep_gemm_supported + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda(), reason="DeepGEMM MegaMoE requires CUDA" +) +requires_sm100 = pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="MegaMoE input staging only feeds the SM100 kernels", +) + + +@pytest.fixture(autouse=True) +def sm100_mega_moe_layout(monkeypatch): + """Build the SM100 kernel layout on any GPU; kernels are gated per test.""" + monkeypatch.setattr( + "vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe." + "check_deep_gemm_mega_moe_supported", + lambda *args: None, + ) + + +def _make_experts(**kwargs) -> DeepGemmMegaMoEExperts: + vllm_config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=4), + compilation_config=SimpleNamespace(static_forward_context={}), + ) + kwargs.setdefault("activation", MegaMoEActivation()) + kwargs.setdefault("weight_format", MEGA_MOE_MXFP4) + return DeepGemmMegaMoEExperts(vllm_config, **kwargs) + + +def test_deep_gemm_mega_moe_capture_precedes_eplb(monkeypatch): + experts = DeepGemmMegaMoEExperts.__new__(DeepGemmMegaMoEExperts) + torch.nn.Module.__init__(experts) + experts._synchronize_first_launch = lambda: None + experts.prefix = "model.layers.3.ffn.experts" + experts.max_num_tokens = 4 + experts.num_shared_experts = 0 + experts.capture_fn = None + experts.eplb_state = SimpleNamespace( + logical_to_physical_map=torch.empty(1), + expert_load_view=torch.empty(1), + logical_replica_count=torch.empty(1), + should_record_tensor=torch.empty(1), + num_unpadded_tokens_tensors=None, + ) + + topk_ids = torch.tensor([[1, 2], [3, 4]]) + captured: list[tuple[int, torch.Tensor]] = [] + bind_routed_experts_capturer( + SimpleNamespace(modules=lambda: [experts]), + SimpleNamespace(capture=lambda layer_id, ids: captured.append((layer_id, ids))), + ) + + class MappingReached(Exception): + pass + + def map_ids(**kwargs): + assert captured == [(3, topk_ids)] + raise MappingReached + + monkeypatch.setattr( + f"{DeepGemmMegaMoEExperts.__module__}.eplb_map_to_physical_and_record", + map_ids, + ) + + with pytest.raises(MappingReached): + experts(torch.empty(2, 8), torch.empty(2, 2), topk_ids) + + +def test_deep_gemm_mega_moe_weight_loader_uses_ep_expert_ownership(): + experts = _make_experts( + num_experts=4, + num_local_experts=2, + experts_start_idx=2, + top_k=2, + hidden_size=128, + intermediate_size=128, + ) + + nonlocal_weight = torch.ones(128, 64, dtype=torch.uint8) + assert ( + experts.weight_loader( + experts.w13_weight, + nonlocal_weight, + "experts.w13_weight", + shard_id="w1", + expert_id=1, + return_success=True, + ) + is False + ) + + w1 = torch.full((128, 64), 3, dtype=torch.uint8) + w3 = torch.full((128, 64), 7, dtype=torch.uint8) + w2 = torch.full((128, 64), 11, dtype=torch.uint8) + + assert experts.weight_loader( + experts.w13_weight, + w1, + "experts.w13_weight", + shard_id="w1", + expert_id=2, + return_success=True, + ) + assert experts.weight_loader( + experts.w13_weight, + w3, + "experts.w13_weight", + shard_id="w3", + expert_id=2, + return_success=True, + ) + assert experts.weight_loader( + experts.w2_weight, + w2, + "experts.w2_weight", + shard_id="w2", + expert_id=2, + return_success=True, + ) + + 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, weight_format=MEGA_MOE_MXFP4 +): + return _make_experts( + 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, + weight_format=weight_format, + ) + + +def _load_mega_moe_checkpoint_shuffled(experts, num_logical_experts, seed): + hidden, inter = experts.hidden_size, experts.intermediate_size + routed = experts.weight_format.routed + group = routed.scale_group_k + shapes = { + "w13_weight": (num_logical_experts, 2 * inter, hidden // 2), + "w13_weight_scale": (num_logical_experts, 2 * inter, hidden // group), + "w2_weight": (num_logical_experts, hidden, inter // 2), + "w2_weight_scale": (num_logical_experts, hidden, inter // group), + } + generator = torch.Generator().manual_seed(seed) + ckpt = { + name: torch.randint(0, 256, shape, dtype=torch.uint8, generator=generator) + for name, shape in shapes.items() + } + if routed.has_global_scale: + for name in ("w13_weight_scale_2", "w2_weight_scale_2", "w2_input_scale"): + ckpt[name] = torch.rand(num_logical_experts, generator=generator) + 0.5 + 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)) + elif tensor.dim() == 1: + # w1 and w3 share one global scale. + shards.append((name, tensor[expert_id], "w1", expert_id)) + shards.append((name, tensor[expert_id], "w3", 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("weight_scale"): + loaded_weight = loaded_weight.view(routed.scale_dtype) + 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 + names = ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale") + ckpt = {name: ckpt[name][physical] for name in names} + 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 + if s13.shape[-1] * 16 == hidden: + # NVFP4: E4M3 bytes packed four per int32, MN-major. + scales = [ + sf.view(torch.int32).transpose(1, 2).contiguous().transpose(1, 2) + for sf in (s13, s2) + ] + else: + 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.l1_weights, *experts.l2_weights) + for tensor, reference in zip(actual, expected): + assert torch.equal(tensor.view(reference.dtype), reference) + + +@pytest.mark.parametrize("weight_format", [MEGA_MOE_MXFP4, MEGA_MOE_NVFP4]) +def test_deep_gemm_mega_moe_loads_shards_into_kernel_layout(weight_format): + """Shards in any order, including EPLB replicas, land in DeepGEMM's layout.""" + experts = _make_mega_moe_experts(4, 3, weight_format) + 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. + group_k = 4 * weight_format.routed.scale_group_k + assert experts.w13_weight_scale.stride() == (768 * 256 // group_k, 1, 768) + assert experts.w2_weight_scale.stride() == (256 * 384 // group_k, 1, 256) + if weight_format.routed.has_global_scale: + physical = torch.arange(4) % 3 + for name in ("w13_weight_scale_2", "w2_weight_scale_2", "w2_input_scale"): + assert torch.equal(getattr(experts, name), ckpt[name][physical]) + assert experts.l2_activation_scale == ckpt["w2_input_scale"].max().item() + + +@pytest.mark.parametrize("weight_format", [MEGA_MOE_MXFP4, MEGA_MOE_NVFP4]) +def test_deep_gemm_mega_moe_reload_overwrites_kernel_layout_in_place(weight_format): + """A second load must update the live kernel tensors without reallocating.""" + experts = _make_mega_moe_experts(2, 2, weight_format) + _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)) + + +def test_deep_gemm_mega_moe_nvfp4_l2_scale_covers_every_expert(): + """Every rank quantizes L2 inputs with the max input scale over all experts, + including those it does not own, so EPLB can move experts between ranks.""" + experts = _make_experts( + num_experts=4, + num_local_experts=2, + experts_start_idx=2, + top_k=2, + hidden_size=128, + intermediate_size=128, + weight_format=MEGA_MOE_NVFP4, + ) + assert experts.l2_activation_scale == 1.0 + for expert_id, scale in enumerate((0.5, 4.0, 1.0, 2.0)): + loaded = experts.weight_loader( + experts.w2_input_scale, + torch.tensor(scale), + "experts.w2_input_scale", + shard_id="w2", + expert_id=expert_id, + return_success=True, + ) + assert loaded == (expert_id >= 2) + assert experts.l2_activation_scale == 4.0 + assert experts.w2_input_scale.tolist() == [1.0, 2.0] + # A reload replaces the expert's scale rather than keeping the old max. + experts.weight_loader( + experts.w2_input_scale, + torch.tensor(0.25), + "experts.w2_input_scale", + shard_id="w2", + expert_id=1, + ) + assert experts.l2_activation_scale == 2.0 + # Global scales are expert weights that EPLB moves. + assert len(experts.get_expert_weights()) == 7 + + +def test_deep_gemm_mega_moe_nvfp4_rejects_mismatched_l1_scales(): + """The kernel takes one L1 alpha per expert, so w1 and w3 must agree.""" + experts = _make_mega_moe_experts(1, 1, MEGA_MOE_NVFP4) + + def load(shard_id, scale): + experts.weight_loader( + experts.w13_weight_scale_2, + torch.tensor(scale), + "experts.w13_weight_scale_2", + shard_id=shard_id, + expert_id=0, + ) + + load("w1", 0.5) + load("w3", 0.5) + # A reload may change the scale while the old pair is still live. + load("w3", 2.0) + load("w1", 2.0) + load("w1", 1.0) + with pytest.raises(ValueError, match="one L1 scale per expert"): + load("w3", 4.0) + + +def test_deep_gemm_mega_moe_nvfp4_meta_load_skips_host_scales(): + """Reload dry runs load meta stand-ins, which must not reach host state.""" + experts = _make_mega_moe_experts(1, 1, MEGA_MOE_NVFP4) + for name, shard_id in ( + ("w2_input_scale", "w2"), + ("w13_weight_scale_2", "w1"), + ("w13_weight_scale_2", "w3"), + ): + assert experts.weight_loader( + getattr(experts, name).to("meta"), + torch.empty((), device="meta"), + f"experts.{name}", + shard_id=shard_id, + expert_id=0, + return_success=True, + ) + assert experts.l2_activation_scale == 1.0 + + +@pytest.mark.parametrize("intermediate_size", [512, 2304]) +def test_deep_gemm_mega_moe_preserves_checkpoint_dimensions(intermediate_size): + """Keep native widths so V4.1 routed and shared experts can fuse.""" + experts = _make_experts( + num_experts=2, + num_local_experts=2, + experts_start_idx=0, + top_k=1, + hidden_size=128, + intermediate_size=intermediate_size, + ) + + 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.l1_weights, *experts.l2_weights) + for actual, param in zip(transformed, experts.parameters()): + assert actual.data_ptr() == param.data_ptr() + + +def _shared_mlp( + hidden_size, + shared_size, + weight_dtype=torch.float8_e4m3fn, + scale_dtype=torch.float8_e8m0fnu, + scale_block=(1, 32), +): + def linear(n, k): + scale = None + if scale_dtype is not None: + scale = torch.empty(n // scale_block[0], k // scale_block[1]).to( + scale_dtype + ) + return SimpleNamespace( + weight=torch.empty(n, k, dtype=weight_dtype), weight_scale_inv=scale + ) + + return SimpleNamespace( + gate_up_proj=linear(2 * shared_size, hidden_size), + down_proj=linear(hidden_size, shared_size), + ) + + +def test_deep_gemm_mega_moe_loads_shared_experts_into_kernel_layout(monkeypatch): + """Shared expert shards land directly in the MegaMoE kernel layout.""" + monkeypatch.setattr( + DeepGemmMegaMoEExperts, "_kernel_fuses_shared_experts", lambda self, spec: True + ) + experts = _make_experts( + num_experts=2, + num_local_experts=2, + experts_start_idx=0, + top_k=1, + hidden_size=256, + intermediate_size=384, + shared_experts=_shared_mlp(256, 384), + ) + 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, shared_dtype=None): ... + + def new_kernel(y, shared_l1_weights=None, shared_l2_weights=None): ... + + kernel = new_kernel if with_shared else old_kernel + buffer = new_buffer if with_shared else old_buffer + return SimpleNamespace( + get_symm_buffer_for_mega_moe=buffer, + NVFP4SymmBuffer=buffer, + fp8_fp4_mega_moe=kernel, + bf16_mega_moe=kernel, + nvfp4_mega_moe=kernel, + get_block_m_for_mega_moe=lambda *args: 0, + get_block_m_for_nvfp4_mega_moe=lambda *args: 0, + ) + + +@pytest.mark.parametrize( + "weight_format,weight_dtype,scale_dtype,scale_block,shared_size,api,fused", + [ + ( + MEGA_MOE_MXFP4, + torch.float8_e4m3fn, + torch.float8_e8m0fnu, + (128, 128), + 128, + True, + True, + ), + (MEGA_MOE_MXFP4, torch.float8_e4m3fn, torch.uint8, (1, 32), 256, True, True), + ( + MEGA_MOE_MXFP4, + torch.float8_e4m3fn, + torch.float8_e8m0fnu, + (128, 128), + 128, + False, + False, + ), + (MEGA_MOE_MXFP4, torch.bfloat16, None, None, 128, True, False), + ( + MEGA_MOE_MXFP4, + torch.float8_e4m3fn, + torch.float32, + (128, 128), + 128, + True, + False, + ), + (MEGA_MOE_MXFP4, torch.float8_e4m3fn, torch.uint8, (1, 16), 128, True, False), + # A TP-sharded shared MLP is not a whole multiple of the routed expert. + (MEGA_MOE_MXFP4, torch.float8_e4m3fn, torch.uint8, (1, 32), 64, True, False), + (MEGA_MOE_BF16, torch.bfloat16, None, None, 128, True, True), + (MEGA_MOE_BF16, torch.float8_e4m3fn, torch.uint8, (1, 32), 128, True, False), + (MEGA_MOE_NVFP4, torch.float8_e4m3fn, torch.uint8, (1, 32), 256, True, True), + (MEGA_MOE_NVFP4, torch.bfloat16, None, None, 128, True, True), + (MEGA_MOE_NVFP4, torch.bfloat16, None, None, 128, False, False), + # NVFP4-style E4M3 block scales are not MXFP8. + ( + MEGA_MOE_NVFP4, + torch.float8_e4m3fn, + torch.float8_e4m3fn, + (1, 16), + 128, + True, + False, + ), + ], +) +def test_deep_gemm_mega_moe_decides_shared_fusion_at_construction( + monkeypatch, + weight_format, + weight_dtype, + scale_dtype, + scale_block, + shared_size, + api, + fused, +): + """Fuse only shared experts in the format's shared precision, shaped as whole + routed experts, and only when the installed kernel takes them.""" + monkeypatch.setattr( + "vllm.utils.deep_gemm._import_deep_gemm", lambda: _fake_deep_gemm(api) + ) + experts = _make_experts( + num_experts=2, + num_local_experts=2, + experts_start_idx=0, + top_k=1, + hidden_size=256, + intermediate_size=128, + weight_format=weight_format, + shared_experts=_shared_mlp( + 256, shared_size, weight_dtype, scale_dtype, scale_block + ), + ) + + assert experts.num_shared_experts == (shared_size // 128 if fused else 0) + assert hasattr(experts, "shared_l1_weight") == fused + if fused: + assert experts.shared_spec in weight_format.shared + assert experts.shared_l1_weight.dtype == weight_dtype + + +def test_deep_gemm_mega_moe_rejects_unsupported_activation(): + with pytest.raises(NotImplementedError, match="situ"): + _make_experts( + num_experts=2, + num_local_experts=2, + experts_start_idx=0, + top_k=1, + hidden_size=128, + intermediate_size=128, + weight_format=MEGA_MOE_BF16, + activation=MegaMoEActivation("situ"), + ) + with pytest.raises(ValueError, match="SiTU takes no clamp"): + MegaMoEActivation("situ", clamp=10.0) + + +@pytest.mark.parametrize("fused", [False, True]) +def test_deep_gemm_mega_moe_adds_unfused_shared_expert_once(monkeypatch, fused): + """The fusion decided at init fixes the launch; forward adds only an + unfused shared expert, and refuses one the layer already fused.""" + monkeypatch.setattr( + DeepGemmMegaMoEExperts, "_kernel_fuses_shared_experts", lambda self, spec: True + ) + launches = [] + + def kernel(y, l1, l2, symm_buffer, **kwargs): + launches.append(kwargs) + y.fill_(1) + + monkeypatch.setattr( + "vllm.utils.deep_gemm._import_deep_gemm", + lambda: SimpleNamespace( + fp8_fp4_mega_moe=kernel, get_block_m_for_mega_moe=lambda *args: 8 + ), + ) + module = "vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe" + monkeypatch.setattr(f"{module}.prepare_megamoe_inputs", lambda *a, **k: None) + monkeypatch.setattr(f"{module}.get_ep_group", lambda: SimpleNamespace(world_size=1)) + shared = _shared_mlp(128, 128 if fused else 64) + experts = _make_experts( + num_experts=2, + num_local_experts=2, + experts_start_idx=0, + top_k=1, + hidden_size=128, + intermediate_size=128, + shared_experts=shared, + ) + assert experts.has_fused_shared_experts == fused + symm_buffer = SimpleNamespace( + x=torch.empty(4, 128), + x_sf=torch.empty(4, 1), + topk_idx=torch.empty(4, 1), + topk_weights=torch.empty(4, 1), + shared_l1_acts_sf=torch.empty(128, 1), + num_max_tokens_per_rank=4, + ) + experts.get_symm_buffer = lambda: symm_buffer + experts._synchronize_first_launch = lambda: None + + def shared_experts(x): + return torch.full_like(x, 2) + + hidden_states = torch.zeros(2, 128, dtype=torch.bfloat16) + topk = (torch.ones(2, 1), torch.zeros(2, 1, dtype=torch.int64)) + if fused: + with pytest.raises(ValueError, match="pass shared_experts=None"): + experts(hidden_states, *topk, shared_experts) + output = experts(hidden_states, *topk) + else: + output = experts(hidden_states, *topk, shared_experts) + + assert torch.all(output == (1 if fused else 3)) + assert ("shared_l1_weights" in launches[-1]) == fused + + +@requires_sm100 +@pytest.mark.parametrize("num_tokens", [7, 63, 64, 65, 127, 128, 135, 434, 16384]) +@pytest.mark.parametrize("hidden_size", [256, 5120]) +def test_prepare_megamoe_inputs_is_bitwise_exact(num_tokens, hidden_size): + from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8 + + device = torch.device("cuda") + top_k = 8 + + generator = torch.Generator(device=device) + generator.manual_seed(0) + hidden_states = ( + torch.randn( + num_tokens, + hidden_size, + device=device, + dtype=torch.float32, + generator=generator, + ) + * 17.0 + ).to(torch.bfloat16) + hidden_states[0, :32] = 0 + hidden_states[1, 32:64] = 1.0e-6 + hidden_states[2, 64:96] = -1.0e-6 + + topk_ids = torch.randint( + 0, + 256, + (num_tokens, top_k), + device=device, + dtype=torch.int32, + generator=generator, + ) + topk_weights = torch.randn( + num_tokens, + top_k, + device=device, + dtype=torch.float32, + generator=generator, + ) + + ref_x, ref_x_sf = per_token_cast_to_fp8( + hidden_states, + use_ue8m0=True, + gran_k=32, + use_packed_ue8m0=True, + ) + ref_topk_idx = topk_ids.to(torch.int64) + ref_topk_weights = topk_weights.clone() + + fused_x = torch.empty_like(ref_x) + fused_x_sf = torch.empty_like(ref_x_sf) + fused_topk_idx = torch.empty_like(ref_topk_idx) + fused_topk_weights = torch.empty_like(ref_topk_weights) + + prepare_megamoe_inputs( + hidden_states, + topk_weights, + topk_ids, + fused_x, + fused_x_sf, + fused_topk_idx, + fused_topk_weights, + ) + torch.accelerator.synchronize() + + assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8)) + assert torch.equal(fused_x_sf, ref_x_sf) + assert torch.equal(fused_topk_idx, ref_topk_idx) + assert torch.equal( + fused_topk_weights.view(torch.uint8), + ref_topk_weights.view(torch.uint8), + ) + + +def _shared_sf_tma_layout(sf: torch.Tensor, block_m: int) -> torch.Tensor: + """Scatter per-token scales into DeepGEMM's padded MN-major UTCCP pages.""" + aligned_block_m = (block_m + 127) // 128 * 128 + num_rows = (sf.shape[0] + block_m - 1) // block_m * aligned_block_m + m = torch.arange(sf.shape[0], device=sf.device) + local = m % block_m + rows = ( + m // block_m * aligned_block_m + + local // 128 * 128 + + local % 32 * 4 + + local % 128 // 32 + ) + out = sf.new_zeros(num_rows, sf.shape[1]) + out[rows] = sf + return out + + +@requires_sm100 +@pytest.mark.parametrize("shared_block_m", [8, 32, 96, 128, 192]) +def test_prepare_megamoe_inputs_stages_shared_scale_tma_layout(shared_block_m): + from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8 + + device = torch.device("cuda") + num_tokens = shared_block_m + 7 + hidden_size = 256 + top_k = 8 + generator = torch.Generator(device=device) + generator.manual_seed(shared_block_m) + hidden_states = torch.randn( + num_tokens, + hidden_size, + device=device, + dtype=torch.bfloat16, + generator=generator, + ) + topk_ids = torch.randint( + 0, + 256, + (num_tokens, top_k), + device=device, + dtype=torch.int32, + generator=generator, + ) + topk_weights = torch.randn( + num_tokens, + top_k, + device=device, + dtype=torch.float32, + generator=generator, + ) + + ref_x, ref_x_sf = per_token_cast_to_fp8( + hidden_states, + use_ue8m0=True, + gran_k=32, + use_packed_ue8m0=True, + ) + ref_shared_x_sf = _shared_sf_tma_layout(ref_x_sf, shared_block_m) + num_shared_rows = ref_shared_x_sf.shape[0] + + fused_x = torch.empty_like(ref_x) + fused_x_sf = torch.empty_like(ref_x_sf) + fused_shared_storage = torch.full( + (hidden_size // 128, num_shared_rows), + -1, + dtype=torch.int32, + device=device, + ) + fused_shared_x_sf = fused_shared_storage.t() + fused_topk_idx = torch.empty_like(topk_ids, dtype=torch.int64) + fused_topk_weights = torch.empty_like(topk_weights) + + prepare_megamoe_inputs( + hidden_states, + topk_weights, + topk_ids, + fused_x, + fused_x_sf, + fused_topk_idx, + fused_topk_weights, + shared_x_sf=fused_shared_x_sf, + shared_block_m=shared_block_m, + ) + torch.accelerator.synchronize() + + populated = ref_shared_x_sf != 0 + assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8)) + assert torch.equal(fused_x_sf, ref_x_sf) + assert torch.equal(fused_shared_x_sf[populated], ref_shared_x_sf[populated]) + assert torch.all(fused_shared_x_sf[~populated] == -1) + + +@requires_sm100 +@pytest.mark.parametrize("quantize", [True, False]) +@pytest.mark.parametrize("nonfinite_padding", [False, True]) +@pytest.mark.parametrize("num_tokens", [7, 65, 135]) +def test_prepare_megamoe_inputs_masks_padding(nonfinite_padding, num_tokens, quantize): + """Padded tokens are unrouted; BF16 dispatch copies the hidden states.""" + from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8 + + device = torch.device("cuda") + hidden_size = 256 + top_k = 8 + + generator = torch.Generator(device=device) + generator.manual_seed(1) + hidden_states = torch.randn( + num_tokens, + hidden_size, + device=device, + dtype=torch.bfloat16, + generator=generator, + ) + topk_ids = torch.randint( + 0, + 256, + (num_tokens, top_k), + device=device, + dtype=torch.int32, + generator=generator, + ) + topk_weights = torch.randn( + num_tokens, + top_k, + device=device, + dtype=torch.float32, + generator=generator, + ) + is_padding = torch.zeros(num_tokens, device=device, dtype=torch.bool) + is_padding[1::3] = True + is_padding[-1] = True + if nonfinite_padding: + topk_weights[is_padding] = float("nan") + + if quantize: + ref_x, ref_x_sf = per_token_cast_to_fp8( + hidden_states, + use_ue8m0=True, + gran_k=32, + use_packed_ue8m0=True, + ) + else: + ref_x, ref_x_sf = hidden_states, None + ref_topk_idx = topk_ids.to(torch.int64) + ref_topk_idx[is_padding] = -1 + ref_topk_weights = topk_weights.clone() + ref_topk_weights[is_padding] = 0.0 + + fused_x = torch.empty_like(ref_x) + fused_x_sf = None if ref_x_sf is None else torch.empty_like(ref_x_sf) + fused_topk_idx = torch.empty_like(ref_topk_idx) + fused_topk_weights = torch.empty_like(ref_topk_weights) + + prepare_megamoe_inputs( + hidden_states, + topk_weights, + topk_ids, + fused_x, + fused_x_sf, + fused_topk_idx, + fused_topk_weights, + is_padding=is_padding, + hidden_quant=kMxfp8Dynamic if quantize else None, + ) + torch.accelerator.synchronize() + + assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8)) + if quantize: + assert torch.equal(fused_x_sf, ref_x_sf) + assert torch.equal(fused_topk_idx, ref_topk_idx) + assert torch.equal( + fused_topk_weights.view(torch.uint8), + ref_topk_weights.view(torch.uint8), + ) + + +@requires_sm100 +@pytest.mark.parametrize("shared", [None, "mxfp8", "bf16"]) +@pytest.mark.parametrize("num_tokens", [0, 7, 65, 135]) +def test_prepare_megamoe_nvfp4_inputs_is_bitwise_exact(num_tokens, shared): + """Per-token NVFP4 dispatch, shared-expert inputs, unrouted padding and the + L2 alphas match DeepGEMM's reference quantizers bit for bit.""" + from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8 + from vllm.third_party.deep_gemm.utils.nvfp4 import per_token_cast_to_nvfp4 + + device = torch.device("cuda") + hidden_size, top_k, shared_block_m = 5120, 8, 32 + generator = torch.Generator(device=device).manual_seed(num_tokens) + # Rows span four decades, which one tensor-wide scale cannot cover. + row_scale = torch.logspace(-3, 1, max(num_tokens, 1), device=device) + hidden_states = ( + torch.randn(num_tokens, hidden_size, device=device, generator=generator) + * row_scale[:num_tokens, None] + ).bfloat16() + if num_tokens: + hidden_states[0] = 0 + topk_ids = torch.randint( + 0, 256, (num_tokens, top_k), device=device, generator=generator + ).int() + topk_weights = torch.rand(num_tokens, top_k, device=device, generator=generator) + is_padding = torch.zeros(num_tokens, device=device, dtype=torch.bool) + is_padding[1::3] = True + w2_global_scale = torch.rand(6, device=device, generator=generator) + + x = torch.empty(num_tokens, hidden_size // 2, dtype=torch.int8, device=device) + x_sf = torch.empty(num_tokens, hidden_size // 64, dtype=torch.int32, device=device) + x_scales = torch.empty(num_tokens, device=device) + topk_idx_out = torch.empty(num_tokens, top_k, dtype=torch.int64, device=device) + topk_weights_out = torch.empty_like(topk_weights) + l2_alpha = torch.empty_like(w2_global_scale) + shared_x = shared_x_sf = None + if shared is not None: + dtype = torch.float8_e4m3fn if shared == "mxfp8" else torch.bfloat16 + shared_x = torch.empty(num_tokens, hidden_size, dtype=dtype, device=device) + if shared == "mxfp8": + ref_shared_x, ref_shared_x_sf = per_token_cast_to_fp8( + hidden_states, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True + ) + ref_shared_x_sf = _shared_sf_tma_layout(ref_shared_x_sf, shared_block_m) + shared_x_sf = torch.full_like(ref_shared_x_sf.t(), -1).t() + + prepare_megamoe_nvfp4_inputs( + hidden_states, + topk_weights, + topk_ids, + x, + x_sf, + x_scales, + topk_idx_out, + topk_weights_out, + is_padding=is_padding, + shared_x=shared_x, + shared_x_sf=shared_x_sf, + shared_block_m=shared_block_m if shared == "mxfp8" else None, + w2_global_scale=w2_global_scale, + l2_alpha=l2_alpha, + l2_activation_scale=0.37, + ) + torch.accelerator.synchronize() + + # Written even without local tokens: other ranks route to these experts. + assert torch.equal(l2_alpha, w2_global_scale * 0.37) + ref_x_scales = hidden_states.float().abs().amax(-1).clamp_min(1e-12) * (1 / 2688) + ref_x, ref_x_sf = per_token_cast_to_nvfp4( + hidden_states.float() / ref_x_scales[:, None] + ) + assert torch.equal(x_scales, ref_x_scales) + assert torch.equal(x, ref_x) + assert torch.equal(x_sf, ref_x_sf) + ref_topk_idx = topk_ids.long().masked_fill(is_padding[:, None], -1) + assert torch.equal(topk_idx_out, ref_topk_idx) + ref_weights = topk_weights.masked_fill(is_padding[:, None], 0.0) + assert torch.equal(topk_weights_out, ref_weights) + if shared == "bf16": + assert torch.equal(shared_x, hidden_states) + elif shared == "mxfp8": + assert shared_x is not None and shared_x_sf is not None + assert torch.equal(shared_x.view(torch.uint8), ref_shared_x.view(torch.uint8)) + populated = ref_shared_x_sf != 0 + assert torch.equal(shared_x_sf[populated], ref_shared_x_sf[populated]) + assert torch.all(shared_x_sf[~populated] == -1) + + +NUM_EXPERTS, TOP_K, HIDDEN, INTER, NUM_TOKENS = 8, 2, 512, 256, 48 + +ACTIVATIONS = { + "swiglu": MegaMoEActivation(), + "swiglu_clamp": MegaMoEActivation(clamp=0.5), + "swiglu_alpha_beta": MegaMoEActivation(alpha=1.702, beta=1.0), + "situ": MegaMoEActivation("situ", alpha=2.0, beta=3.0), +} + + +def _quantize(w: torch.Tensor, spec: MegaMoEWeightSpec): + """Return checkpoint ``(weight, scale, global_scale)`` and the dequantized + weight; NVFP4 takes ``global_scale`` as ``weight_scale_2``.""" + from vllm.third_party.deep_gemm.utils.math import ( + cast_back_from_fp4, + cast_back_from_fp8, + per_token_cast_to_fp4, + per_token_cast_to_fp8, + ) + from vllm.third_party.deep_gemm.utils.nvfp4 import ( + cast_back_from_nvfp4, + per_token_cast_to_nvfp4, + ) + + if spec.scale_group_k is None: + return w, None, None, w.float() + if spec.has_global_scale: + global_scale = w.float().abs().amax().item() / 2688 + q, sf = per_token_cast_to_nvfp4(w, global_scale) + deq = cast_back_from_nvfp4(q, sf, global_scale) + return q, sf.view(torch.float8_e4m3fn), torch.tensor(global_scale), deq + if spec.values_per_element == 2: + q, sf = per_token_cast_to_fp4(w, use_ue8m0=True, gran_k=32) + deq = cast_back_from_fp4(q, sf, gran_k=32) + else: + q, sf = per_token_cast_to_fp8(w, use_ue8m0=True, gran_k=32) + deq = cast_back_from_fp8(q, sf, gran_k=32) + return q, (sf.view(torch.int32) >> 23).to(torch.uint8), None, deq + + +def _round_trip(x: torch.Tensor, quant, nvfp4_scale: float | None = None): + """What the kernel computes on: x after ``quant`` (None keeps BF16).""" + from vllm.third_party.deep_gemm.utils.math import ( + cast_back_from_fp8, + per_token_cast_to_fp8, + ) + from vllm.third_party.deep_gemm.utils.nvfp4 import ( + cast_back_from_nvfp4, + per_token_cast_to_nvfp4, + ) + + if quant is None: + return x.bfloat16().float() + if quant == kNvfp4DynamicToken: + if nvfp4_scale is not None: + q = per_token_cast_to_nvfp4(x, nvfp4_scale) + return cast_back_from_nvfp4(*q, nvfp4_scale) + x = x.bfloat16().float() + scales = x.abs().amax(-1, keepdim=True).clamp_min(1e-12) / 2688 + return cast_back_from_nvfp4(*per_token_cast_to_nvfp4(x / scales)) * scales + return cast_back_from_fp8( + *per_token_cast_to_fp8(x.bfloat16(), use_ue8m0=True, gran_k=32), gran_k=32 + ) + + +def _activate(h: torch.Tensor, act: MegaMoEActivation) -> torch.Tensor: + gate, up = h.bfloat16().float().chunk(2, dim=-1) + if act.name == "situ": + gate = act.alpha * torch.tanh(gate / act.alpha) * torch.sigmoid(gate) + up = act.beta * torch.tanh(up / act.beta) if act.beta else up + return gate * up + if act.clamp is not None: + gate, up = gate.clamp(max=act.clamp), up.clamp(-act.clamp, act.clamp) + return gate * torch.sigmoid(act.alpha * gate) * (up + act.beta) + + +def _expert(x, w13, w2, weight, fmt, act, l2_scale): + """Reference routed expert. MX and BF16 weight the activation before the + L2 dispatch; NVFP4 quantizes the unweighted activation with the global + ``l2_scale`` and weights the BF16 output.""" + a = _activate(x @ w13.t(), act) + if fmt.hidden_quant == kNvfp4DynamicToken: + a = _round_trip(a, fmt.hidden_quant, l2_scale) + return (a @ w2.t()).bfloat16().float() * weight[:, None] + return _round_trip(a * weight[:, None], fmt.hidden_quant) @ w2.t() + + +class _ReferenceSharedExpert(torch.nn.Module): + def __init__(self, w13, w2, spec: MegaMoEWeightSpec, act): + super().__init__() + self.w13, self.w2, self.act = w13, w2, act + self.quant = kMxfp8Dynamic if spec.scale_group_k else None + + def forward(self, x): + a = _activate(_round_trip(x, self.quant) @ self.w13.t(), self.act) + return (_round_trip(a, self.quant) @ self.w2.t()).bfloat16() + + +def _calc_diff(x: torch.Tensor, y: torch.Tensor) -> float: + x, y = x.double(), y.double() + return (1 - 2 * (x * y).sum() / (x * x + y * y).sum()).item() + + +@pytest.mark.skipif( + not ( + current_platform.is_device_capability_family(100) and is_deep_gemm_supported() + ), + reason="DeepGEMM MegaMoE kernels require SM100 and DeepGEMM", +) +@pytest.mark.parametrize( + "fmt_name,shared,act_name", + [ + pytest.param( + name, + shared, + act_name, + id=f"{act_name}-{name}-{str(shared).removeprefix('torch.')}", + ) + for act_name, act in ACTIVATIONS.items() + for name, fmt in MEGA_MOE_WEIGHT_FORMATS.items() + if act.name in fmt.activations + for shared in (None, *(spec.dtype for spec in fmt.shared)) + ], +) +def test_deep_gemm_mega_moe_matches_reference(dist_init, fmt_name, shared, act_name): + """``shared`` is the fused shared-expert dtype; None runs it unfused.""" + fmt, act = MEGA_MOE_WEIGHT_FORMATS[fmt_name], ACTIVATIONS[act_name] + shared_spec = next( + (spec for spec in fmt.shared if spec.dtype == shared), fmt.shared[0] + ) + fuse_shared = shared is not None + torch.manual_seed(0) + device = torch.device("cuda") + vllm_config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=NUM_TOKENS), + compilation_config=SimpleNamespace(static_forward_context={}), + ) + + def randn(*shape, fan_in): + return (torch.randn(*shape, device=device) / fan_in**0.5).bfloat16() + + shared_ckpt = [ + _quantize(randn(*shape, fan_in=k), shared_spec) + for shape, k in (((INTER, HIDDEN), HIDDEN),) * 2 + (((HIDDEN, INTER), INTER),) + ] + shared_mlp = SimpleNamespace( + gate_up_proj=SimpleNamespace( + weight=torch.cat([shared_ckpt[0][0], shared_ckpt[1][0]]), + weight_scale_inv=None + if shared_ckpt[0][1] is None + else torch.cat([shared_ckpt[0][1], shared_ckpt[1][1]]), + ), + down_proj=SimpleNamespace( + weight=shared_ckpt[2][0], weight_scale_inv=shared_ckpt[2][1] + ), + ) + + with torch.device(device): + experts = DeepGemmMegaMoEExperts( + vllm_config, + num_experts=NUM_EXPERTS, + num_local_experts=NUM_EXPERTS, + experts_start_idx=0, + top_k=TOP_K, + hidden_size=HIDDEN, + intermediate_size=INTER, + activation=act, + weight_format=fmt, + shared_experts=shared_mlp if fuse_shared else None, + prefix=f"model.layers.0.{fmt_name}.{act_name}.{shared}", + ) + assert experts.has_fused_shared_experts == fuse_shared + assert experts.shared_spec == (shared_spec if fuse_shared else None) + + w13 = torch.empty(NUM_EXPERTS, 2 * INTER, HIDDEN, device=device) + w2 = torch.empty(NUM_EXPERTS, HIDDEN, INTER, device=device) + # NVFP4 L2 input scales; the kernel uses their max. + l2_scales = [2e-3 * (e + 1) for e in range(NUM_EXPERTS)] + for e in range(NUM_EXPERTS): + # w1 and w3 are quantized together: NVFP4 needs one global scale. + w13_ckpt = _quantize(randn(2 * INTER, HIDDEN, fan_in=HIDDEN), fmt.routed) + w2_ckpt = _quantize(randn(HIDDEN, INTER, fan_in=INTER), fmt.routed) + w13[e], w2[e] = w13_ckpt[3], w2_ckpt[3] + for shard_id, param, rows, (weight, scale, global_scale, _) in ( + ("w1", "w13", slice(0, INTER), w13_ckpt), + ("w3", "w13", slice(INTER, None), w13_ckpt), + ("w2", "w2", slice(None), w2_ckpt), + ): + tensors = { + "weight": weight[rows], + "weight_scale": None if scale is None else scale[rows], + "weight_scale_2": global_scale, + } + if shard_id == "w2" and global_scale is not None: + tensors["input_scale"] = torch.tensor(l2_scales[e]) + for suffix, tensor in tensors.items(): + if tensor is None: + continue + name = f"{param}_{suffix}" + assert experts.weight_loader( + getattr(experts, name), + tensor, + f"experts.{name}", + shard_id=shard_id, + expert_id=e, + return_success=True, + ) + + shared_w13 = torch.cat([shared_ckpt[0][3], shared_ckpt[1][3]]) + reference_shared = _ReferenceSharedExpert( + shared_w13, shared_ckpt[2][3], shared_spec, act + ) + if fuse_shared: + for (weight, scale, _, _), shard in zip(shared_ckpt, ("w1", "w3", "w2")): + experts.load_shared_expert_weight(f"shared_experts.{shard}.weight", weight) + if scale is not None: + experts.load_shared_expert_weight( + f"shared_experts.{shard}.weight_scale", scale + ) + + hidden_states = torch.randn(NUM_TOKENS, HIDDEN, device=device).bfloat16() + scores = torch.randn(NUM_TOKENS, NUM_EXPERTS, device=device) + topk_weights, topk_ids = torch.topk(scores.softmax(-1), TOP_K, dim=-1) + + output = experts( + hidden_states, + topk_weights, + topk_ids, + None if fuse_shared else reference_shared, + ) + + x = _round_trip(hidden_states.float(), fmt.hidden_quant) + l2_scale = experts.l2_activation_scale + expected = reference_shared(hidden_states).float() + for e in range(NUM_EXPERTS): + token, slot = (topk_ids == e).nonzero(as_tuple=True) + expected.index_add_( + 0, + token, + _expert( + x[token], w13[e], w2[e], topk_weights[token, slot], fmt, act, l2_scale + ), + ) + + assert output.dtype == torch.bfloat16 + assert _calc_diff(output, expected) < 1e-3 + assert F.cosine_similarity(output.float(), expected, dim=-1).min() > 0.99 diff --git a/tests/model_executor/test_deep_gemm_warmup.py b/tests/model_executor/test_deep_gemm_warmup.py index c2dc12a27d..26c410d1bf 100644 --- a/tests/model_executor/test_deep_gemm_warmup.py +++ b/tests/model_executor/test_deep_gemm_warmup.py @@ -123,19 +123,12 @@ def test_mega_moe_m_values_reach_mega_gate(monkeypatch) -> None: threshold (1 for 128 experts), so the forward actually reaches the kernel.""" dg = SimpleNamespace( get_bf16_mega_gate_config=lambda m, *_: {"num_sms": -(-m // 16)}, - get_block_m_for_mega_moe=lambda *_: 16, ) monkeypatch.setattr(deep_gemm_warmup, "_import_deep_gemm", lambda: dg) monkeypatch.setattr(deep_gemm_warmup, "_mega_moe_uses_mega_gate", lambda _: True) - monkeypatch.setattr( - deep_gemm_warmup, "get_ep_group", lambda: SimpleNamespace(world_size=8) - ) - buffer = SimpleNamespace(num_max_tokens_per_rank=64) module = SimpleNamespace( gate=SimpleNamespace(input_size=4096, output_size=128), - experts=SimpleNamespace( - top_k=6, num_experts=128, get_symm_buffer=lambda: buffer - ), + experts=SimpleNamespace(top_k=6, get_block_m=lambda num_tokens: 16), use_sequence_parallel=False, ) assert deep_gemm_warmup._get_fp8_fp4_mega_moe_m_values(module, 40) == [16, 32, 40] diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index 97dd1b376b..aa6e9c976d 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -8,10 +8,8 @@ import torch from transformers import DeepseekV4Config from vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe import ( - DeepGemmSm100MegaMoEBackend, -) -from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( - bind_routed_experts_capturer, + DeepGemmMegaMoEExperts, + MegaMoEActivation, ) from vllm.model_executor.models.interfaces import ( get_mixture_of_experts_model, @@ -20,17 +18,14 @@ from vllm.model_executor.models.interfaces import ( from vllm.models.deepseek_v4.nvidia.dspark import DSparkDeepseekV4ForCausalLM from vllm.models.deepseek_v4.nvidia.model import ( DeepseekV4ForCausalLM, - DeepseekV4MegaMoEExperts, DeepseekV4MoE, make_deepseek_v4_expert_params_mapping, prepare_mega_gate_routing_metadata, ) -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 from vllm.platforms import current_platform from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config -from vllm.utils.deep_gemm import is_deep_gemm_supported from vllm.utils.torch_utils import set_default_torch_dtype pytestmark = pytest.mark.skipif( @@ -40,11 +35,12 @@ pytestmark = pytest.mark.skipif( @pytest.fixture(autouse=True) -def sm100_mega_moe_backend(monkeypatch): +def sm100_mega_moe_layout(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(), + "vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe." + "check_deep_gemm_mega_moe_supported", + lambda *args: None, ) @@ -125,7 +121,8 @@ def test_deepseek_v41_moe_routes_without_hash_table( ) assert moe.gate.tid2eid is None assert moe.gate.weight.shape == (num_experts, config.hidden_size) - assert isinstance(moe.experts, DeepseekV4MegaMoEExperts) + assert isinstance(moe.experts, DeepGemmMegaMoEExperts) + assert moe.experts.activation == MegaMoEActivation("swiglu", clamp=10.0) assert moe.experts.w13_weight.shape == (num_experts, 256, 128) assert moe.experts.top_k == top_k assert (config.n_routed_experts, config.num_experts_per_tok) == (384, 6) @@ -151,10 +148,9 @@ def test_deepseek_v41_moe_routes_without_hash_table( routed = {} - def check_routing(x, weights, ids, *, activation_clamp): + def check_routing(x, weights, ids, shared_experts): routed["ids"], routed["weights"] = ids, weights - assert activation_clamp == config.swiglu_limit - return x.clone() + return x + shared_experts(x) monkeypatch.setattr(moe.experts, "forward", check_routing) monkeypatch.setattr(moe.shared_experts, "forward", lambda x: 2 * x) @@ -280,10 +276,10 @@ def test_deepseek_v4_mega_gate_hash_routing_correctness( dim=-1, keepdim=True ) - def check_routing(x, weights, ids, *, activation_clamp): + def check_routing(x, weights, ids, shared_experts): torch.testing.assert_close(ids, expected_ids) torch.testing.assert_close(weights, expected_weights, rtol=1e-3, atol=1e-4) - return x.clone() + return x + shared_experts(x) monkeypatch.setattr(moe.experts, "forward", check_routing) monkeypatch.setattr(moe.shared_experts, "forward", torch.zeros_like) @@ -308,634 +304,6 @@ def test_deepseek_v4_mega_moe_expert_mapping(): ] -@pytest.mark.parametrize("use_kimi", [False, True]) -def test_deep_gemm_mega_moe_capture_precedes_eplb(monkeypatch, use_kimi): - experts_cls = DeepseekV4MegaMoEExperts - if use_kimi: - from vllm.models.kimi_k3.nvidia.model import KimiK3MegaMoEExperts - - experts_cls = KimiK3MegaMoEExperts - - experts = experts_cls.__new__(experts_cls) - torch.nn.Module.__init__(experts) - if use_kimi: - experts.synchronize_first_launch = lambda: None - experts.prefix = "model.layers.3.ffn.experts" - experts.max_num_tokens = 4 - experts.capture_fn = None - experts.get_symm_buffer = lambda: object() - experts.eplb_state = SimpleNamespace( - logical_to_physical_map=torch.empty(1), - expert_load_view=torch.empty(1), - logical_replica_count=torch.empty(1), - should_record_tensor=torch.empty(1), - num_unpadded_tokens_tensors=None, - ) - - topk_ids = torch.tensor([[1, 2], [3, 4]]) - captured: list[tuple[int, torch.Tensor]] = [] - bind_routed_experts_capturer( - SimpleNamespace(modules=lambda: [experts]), - SimpleNamespace(capture=lambda layer_id, ids: captured.append((layer_id, ids))), - ) - - class MappingReached(Exception): - pass - - def map_ids(**kwargs): - assert captured == [(3, topk_ids)] - raise MappingReached - - monkeypatch.setattr( - f"{experts_cls.__module__}.eplb_map_to_physical_and_record", - map_ids, - ) - monkeypatch.setattr( - "vllm.utils.deep_gemm._import_deep_gemm", lambda: SimpleNamespace() - ) - - with pytest.raises(MappingReached): - experts( - torch.empty(2, 8), - torch.empty(2, 2), - topk_ids, - activation_clamp=None, - ) - - -def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership(): - vllm_config = SimpleNamespace( - scheduler_config=SimpleNamespace(max_num_batched_tokens=4), - compilation_config=SimpleNamespace(static_forward_context={}), - ) - experts = DeepseekV4MegaMoEExperts( - vllm_config, - num_experts=4, - num_local_experts=2, - experts_start_idx=2, - top_k=2, - hidden_size=128, - intermediate_size=128, - ) - - nonlocal_weight = torch.ones(128, 64, dtype=torch.uint8) - assert ( - experts.weight_loader( - experts.w13_weight, - nonlocal_weight, - "experts.w13_weight", - shard_id="w1", - expert_id=1, - return_success=True, - ) - is False - ) - - w1 = torch.full((128, 64), 3, dtype=torch.uint8) - w3 = torch.full((128, 64), 7, dtype=torch.uint8) - w2 = torch.full((128, 64), 11, dtype=torch.uint8) - - assert experts.weight_loader( - experts.w13_weight, - w1, - "experts.w13_weight", - shard_id="w1", - expert_id=2, - return_success=True, - ) - assert experts.weight_loader( - experts.w13_weight, - w3, - "experts.w13_weight", - shard_id="w3", - expert_id=2, - return_success=True, - ) - assert experts.weight_loader( - experts.w2_weight, - w2, - "experts.w2_weight", - shard_id="w2", - expert_id=2, - return_success=True, - ) - - 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(intermediate_size): - """Keep native widths so V4.1 routed and shared experts can fuse.""" - 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=128, - intermediate_size=intermediate_size, - ) - - 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, param in zip(transformed, experts.parameters()): - assert actual.data_ptr() == param.data_ptr() - - -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( - "weight_dtype,scale_dtype,scale_shape,with_shared,fused", - [ - (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_decides_shared_fusion_at_construction( - monkeypatch, weight_dtype, scale_dtype, scale_shape, with_shared, fused -): - """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( - 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=128, - num_shared_experts=1, - shared_experts=shared, - ) - - assert experts.has_fused_shared_experts == fused - assert hasattr(experts, "shared_l1_weight") == fused - - -@pytest.mark.parametrize("fused", [False, True]) -def test_deepseek_v4_mega_moe_does_not_double_add_fused_shared_expert( - monkeypatch, fused -): - class FakeGate(torch.nn.Module): - weight = torch.empty(2, 128) - tid2eid = None - e_score_correction_bias = None - - def forward(self, hidden_states): - return torch.empty(hidden_states.shape[0], 2), None - - class FakeExperts(torch.nn.Module): - has_fused_shared_experts = fused - backend = DeepGemmSm100MegaMoEBackend() - - def forward(self, hidden_states, *args, **kwargs): - return torch.ones_like(hidden_states) - - class FakeSharedExperts(torch.nn.Module): - calls = 0 - - def forward(self, hidden_states): - self.calls += 1 - return torch.full_like(hidden_states, 2) - - moe = DeepseekV4MoE.__new__(DeepseekV4MoE) - torch.nn.Module.__init__(moe) - moe.use_native_mega_moe = True - moe.gate = FakeGate() - moe.experts = FakeExperts() - moe.shared_experts = None if fused else FakeSharedExperts() - moe.scoring_func = "sqrtsoftplus" - moe.n_activated_experts = 1 - moe.renormalize = True - moe.hash_indices_dtype = torch.int64 - moe.routed_scaling_factor = 1.0 - moe.swiglu_limit = 10.0 - monkeypatch.setattr( - "vllm.models.deepseek_v4.nvidia.model.fused_topk_bias", - lambda **kwargs: ( - torch.ones(kwargs["hidden_states"].shape[0], 1), - torch.zeros(kwargs["hidden_states"].shape[0], 1, dtype=torch.int64), - ), - ) - - output = moe(torch.zeros(2, 128)) - - expected = 1 if fused else 3 - assert torch.all(output == expected) - if not fused: - assert moe.shared_experts.calls == 1 - - -@pytest.mark.skipif( - not torch.cuda.is_available(), - reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.", -) -@pytest.mark.parametrize("num_tokens", [7, 63, 64, 65, 127, 128, 135, 434, 16384]) -@pytest.mark.parametrize("hidden_size", [256, 5120]) -def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact( - num_tokens, hidden_size -): - from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8 - - device = torch.device("cuda") - top_k = 8 - - generator = torch.Generator(device=device) - generator.manual_seed(0) - hidden_states = ( - torch.randn( - num_tokens, - hidden_size, - device=device, - dtype=torch.float32, - generator=generator, - ) - * 17.0 - ).to(torch.bfloat16) - hidden_states[0, :32] = 0 - hidden_states[1, 32:64] = 1.0e-6 - hidden_states[2, 64:96] = -1.0e-6 - - topk_ids = torch.randint( - 0, - 256, - (num_tokens, top_k), - device=device, - dtype=torch.int32, - generator=generator, - ) - topk_weights = torch.randn( - num_tokens, - top_k, - device=device, - dtype=torch.float32, - generator=generator, - ) - - ref_x, ref_x_sf = per_token_cast_to_fp8( - hidden_states, - use_ue8m0=True, - gran_k=32, - use_packed_ue8m0=True, - ) - ref_topk_idx = topk_ids.to(torch.int64) - ref_topk_weights = topk_weights.clone() - - fused_x = torch.empty_like(ref_x) - fused_x_sf = torch.empty_like(ref_x_sf) - fused_topk_idx = torch.empty_like(ref_topk_idx) - fused_topk_weights = torch.empty_like(ref_topk_weights) - - prepare_megamoe_inputs( - hidden_states, - topk_weights, - topk_ids, - fused_x, - fused_x_sf, - fused_topk_idx, - fused_topk_weights, - ) - torch.accelerator.synchronize() - - assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8)) - assert torch.equal(fused_x_sf, ref_x_sf) - assert torch.equal(fused_topk_idx, ref_topk_idx) - assert torch.equal( - fused_topk_weights.view(torch.uint8), - ref_topk_weights.view(torch.uint8), - ) - - -@pytest.mark.parametrize("shared_block_m", [8, 32, 96, 128, 192]) -def test_deepseek_v4_mega_moe_stages_shared_scale_tma_layout(shared_block_m): - from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8 - - device = torch.device("cuda") - num_tokens = shared_block_m + 7 - hidden_size = 256 - top_k = 8 - generator = torch.Generator(device=device) - generator.manual_seed(shared_block_m) - hidden_states = torch.randn( - num_tokens, - hidden_size, - device=device, - dtype=torch.bfloat16, - generator=generator, - ) - topk_ids = torch.randint( - 0, - 256, - (num_tokens, top_k), - device=device, - dtype=torch.int32, - generator=generator, - ) - topk_weights = torch.randn( - num_tokens, - top_k, - device=device, - dtype=torch.float32, - generator=generator, - ) - - ref_x, ref_x_sf = per_token_cast_to_fp8( - hidden_states, - use_ue8m0=True, - gran_k=32, - use_packed_ue8m0=True, - ) - aligned_block_m = ((shared_block_m + 127) // 128) * 128 - num_shared_rows = ((num_tokens + shared_block_m - 1) // shared_block_m) * ( - aligned_block_m - ) - ref_shared_x_sf = torch.zeros( - num_shared_rows, - hidden_size // 128, - dtype=torch.int32, - device=device, - ) - for token_id in range(num_tokens): - m_in_block = token_id % shared_block_m - transposed_m = ( - (m_in_block // 128) * 128 + (m_in_block % 32) * 4 + (m_in_block % 128) // 32 - ) - shared_row = token_id // shared_block_m * aligned_block_m + transposed_m - ref_shared_x_sf[shared_row].copy_(ref_x_sf[token_id]) - - fused_x = torch.empty_like(ref_x) - fused_x_sf = torch.empty_like(ref_x_sf) - fused_shared_storage = torch.full( - (hidden_size // 128, num_shared_rows), - -1, - dtype=torch.int32, - device=device, - ) - fused_shared_x_sf = fused_shared_storage.t() - fused_topk_idx = torch.empty_like(topk_ids, dtype=torch.int64) - fused_topk_weights = torch.empty_like(topk_weights) - - prepare_megamoe_inputs( - hidden_states, - topk_weights, - topk_ids, - fused_x, - fused_x_sf, - fused_topk_idx, - fused_topk_weights, - shared_x_sf=fused_shared_x_sf, - shared_block_m=shared_block_m, - ) - torch.accelerator.synchronize() - - populated = ref_shared_x_sf != 0 - assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8)) - assert torch.equal(fused_x_sf, ref_x_sf) - assert torch.equal(fused_shared_x_sf[populated], ref_shared_x_sf[populated]) - assert torch.all(fused_shared_x_sf[~populated] == -1) - - 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 mhc broadcast weights to cover dummy load, which @@ -1040,85 +408,3 @@ def test_dspark_draft_registers_mixture_of_experts( assert len(draft.moe_layers) == draft.model.num_dspark_layers == 2 assert draft.num_routed_experts == hf_config.n_routed_experts assert draft.num_redundant_experts == 4 - - -@pytest.mark.skipif( - not torch.cuda.is_available(), - reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.", -) -@pytest.mark.parametrize("nonfinite_padding", [False, True]) -@pytest.mark.parametrize("num_tokens", [7, 65, 135]) -def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding( - nonfinite_padding, num_tokens -): - from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8 - - device = torch.device("cuda") - hidden_size = 256 - top_k = 8 - - generator = torch.Generator(device=device) - generator.manual_seed(1) - hidden_states = torch.randn( - num_tokens, - hidden_size, - device=device, - dtype=torch.bfloat16, - generator=generator, - ) - topk_ids = torch.randint( - 0, - 256, - (num_tokens, top_k), - device=device, - dtype=torch.int32, - generator=generator, - ) - topk_weights = torch.randn( - num_tokens, - top_k, - device=device, - dtype=torch.float32, - generator=generator, - ) - is_padding = torch.zeros(num_tokens, device=device, dtype=torch.bool) - is_padding[1::3] = True - is_padding[-1] = True - if nonfinite_padding: - topk_weights[is_padding] = float("nan") - - ref_x, ref_x_sf = per_token_cast_to_fp8( - hidden_states, - use_ue8m0=True, - gran_k=32, - use_packed_ue8m0=True, - ) - ref_topk_idx = topk_ids.to(torch.int64) - ref_topk_idx[is_padding] = -1 - ref_topk_weights = topk_weights.clone() - ref_topk_weights[is_padding] = 0.0 - - fused_x = torch.empty_like(ref_x) - fused_x_sf = torch.empty_like(ref_x_sf) - fused_topk_idx = torch.empty_like(ref_topk_idx) - fused_topk_weights = torch.empty_like(ref_topk_weights) - - prepare_megamoe_inputs( - hidden_states, - topk_weights, - topk_ids, - fused_x, - fused_x_sf, - fused_topk_idx, - fused_topk_weights, - is_padding=is_padding, - ) - torch.accelerator.synchronize() - - assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8)) - assert torch.equal(fused_x_sf, ref_x_sf) - assert torch.equal(fused_topk_idx, ref_topk_idx) - assert torch.equal( - fused_topk_weights.view(torch.uint8), - ref_topk_weights.view(torch.uint8), - ) 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 b86820f672..69a234b081 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 @@ -1,170 +1,559 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""DeepGEMM MegaMoE backends: fused EP dispatch, expert GEMMs and combine.""" +"""DeepGEMM MegaMoE experts: fused EP dispatch, expert GEMMs and combine. +Each ``MEGA_MOE_*`` format is one row of the SM100 feature matrix. MX weights +and activations carry one UE8M0 scale per 1x32 block. NVFP4 weights carry one +E4M3 scale per 1x16 block plus a per-expert FP32 global scale; NVFP4 dispatch +scales each token by its own FP32 scale. +""" + +import math +from collections.abc import Callable +from dataclasses import dataclass from inspect import signature -from typing import Any +from typing import Any, ClassVar, Literal import torch import torch.nn as nn +import vllm.envs as envs +from vllm.config import VllmConfig +from vllm.distributed import get_ep_group +from vllm.distributed.eplb.eplb_state import EplbLayerState +from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.prepare_megamoe import ( + prepare_megamoe_inputs, + prepare_megamoe_nvfp4_inputs, +) +from vllm.model_executor.layers.fused_moe.router.base_router import ( + eplb_map_to_physical_and_record, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kMxfp8Dynamic, + kNvfp4DynamicToken, ) +from vllm.model_executor.models.utils import extract_layer_index +from vllm.model_executor.utils import set_weight_attrs from vllm.platforms import current_platform +from vllm.v1.worker.ubatching import dbo_current_ubatch_id logger = init_logger(__name__) -class DeepGemmMegaMoEBackend: - """Abstract MegaMoE backend that hides deep_gemm-specific details. +@dataclass(frozen=True) +class MegaMoEActivation: + """Gated activation fused into the MegaMoE L1 epilogue. - A backend advertises its expected hidden-state quantization layout via - ``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. + ``swiglu`` computes ``gate * sigmoid(alpha * gate) * (up + beta)`` after + clamping gate from above and up symmetrically to ``clamp``; this covers + plain, DeepSeek-V4 (``clamp``) and GPT-OSS style (``alpha``/``beta``) + experts. ``situ`` is Kimi K3's SiTU with gate tanh scale ``alpha`` and + linear tanh scale ``beta``; ``beta=0`` leaves the linear branch unscaled. """ - # Human-readable MMA type tag used by deep_gemm helpers. SM100 default. - mma_type: str = "fp8xfp4" - hidden_quant: QuantKey = kMxfp8Dynamic + name: Literal["swiglu", "situ"] = "swiglu" + clamp: float | None = None + alpha: float = 1.0 + beta: float = 0.0 - def create_weights( + def __post_init__(self) -> None: + if self.name not in ("swiglu", "situ"): + raise ValueError(f"Unsupported MegaMoE activation {self.name!r}.") + if not (math.isfinite(self.alpha) and math.isfinite(self.beta)): + raise ValueError(f"MegaMoE activation scales must be finite: {self}.") + if self.clamp is not None and not self.clamp >= 0: + raise ValueError(f"MegaMoE activation clamp must be >= 0: {self}.") + if self.name == "situ" and ( + self.clamp is not None or self.alpha <= 0 or self.beta < 0 + ): + raise ValueError( + "SiTU takes no clamp, a positive alpha and a non-negative beta; " + f"got {self}." + ) + + def kernel_kwargs(self) -> dict[str, Any]: + # Defaults stay implicit so DeepGEMM builds without alpha/beta still work. + kwargs: dict[str, Any] = {"activation_clamp": self.clamp} + if self.name != "swiglu": + kwargs["activation"] = self.name + if self.alpha != 1.0: + kwargs["activation_alpha"] = self.alpha + if self.beta != 0.0: + kwargs["activation_beta"] = self.beta + return kwargs + + +@dataclass(frozen=True) +class MegaMoEWeightSpec: + """Storage of one MegaMoE weight in the kernel layout.""" + + dtype: torch.dtype + # Weight values packed per storage element along K (2 for FP4). + values_per_element: int = 1 + # K values per block scale; None for unscaled weights. + scale_group_k: int | None = None + # Block scale dtype: UE8M0 (MX) or E4M3 (NVFP4). + scale_dtype: torch.dtype | None = None + + @property + def has_global_scale(self) -> bool: + """E4M3 block scales pair with an FP32 per-tensor scale.""" + return self.scale_dtype == torch.float8_e4m3fn + + +_MXFP4_WEIGHT = MegaMoEWeightSpec( + torch.int8, + values_per_element=2, + scale_group_k=32, + scale_dtype=torch.float8_e8m0fnu, +) +_MXFP8_WEIGHT = MegaMoEWeightSpec( + torch.float8_e4m3fn, scale_group_k=32, scale_dtype=torch.float8_e8m0fnu +) +_NVFP4_WEIGHT = MegaMoEWeightSpec( + torch.int8, + values_per_element=2, + scale_group_k=16, + scale_dtype=torch.float8_e4m3fn, +) +_BF16_WEIGHT = MegaMoEWeightSpec(torch.bfloat16) + + +@dataclass(frozen=True) +class MegaMoEWeightFormat: + """One routed-expert precision of the MegaMoE feature matrix.""" + + name: str + # DeepGEMM MMA tag; it also keys the symmetric-buffer layout. + mma_type: str + kernel: str + routed: MegaMoEWeightSpec + # Precisions a shared expert may have to run inside the same launch. + shared: tuple[MegaMoEWeightSpec, ...] + # Quantization of the dispatched hidden states; None dispatches BF16. + hidden_quant: QuantKey | None + activations: frozenset[str] + + @property + def uses_nvfp4_api(self) -> bool: + """``nvfp4_mega_moe`` has its own symmetric buffer, BLOCK_M and scales.""" + return self.kernel == "nvfp4_mega_moe" + + +MEGA_MOE_MXFP4 = MegaMoEWeightFormat( + name="mxfp4", + mma_type="fp8xfp4", + kernel="fp8_fp4_mega_moe", + routed=_MXFP4_WEIGHT, + shared=(_MXFP8_WEIGHT,), + hidden_quant=kMxfp8Dynamic, + activations=frozenset({"swiglu", "situ"}), +) +MEGA_MOE_MXFP8 = MegaMoEWeightFormat( + name="mxfp8", + mma_type="fp8xfp8", + kernel="fp8_fp4_mega_moe", + routed=_MXFP8_WEIGHT, + shared=(_MXFP8_WEIGHT,), + hidden_quant=kMxfp8Dynamic, + activations=frozenset({"swiglu", "situ"}), +) +MEGA_MOE_NVFP4 = MegaMoEWeightFormat( + name="nvfp4", + mma_type="nvfp4", + kernel="nvfp4_mega_moe", + routed=_NVFP4_WEIGHT, + shared=(_MXFP8_WEIGHT, _BF16_WEIGHT), + hidden_quant=kNvfp4DynamicToken, + activations=frozenset({"swiglu"}), +) +MEGA_MOE_BF16 = MegaMoEWeightFormat( + name="bf16", + mma_type="bf16xbf16", + kernel="bf16_mega_moe", + routed=_BF16_WEIGHT, + shared=(_BF16_WEIGHT,), + hidden_quant=None, + activations=frozenset({"swiglu"}), +) +MEGA_MOE_WEIGHT_FORMATS = { + fmt.name: fmt + for fmt in (MEGA_MOE_MXFP4, MEGA_MOE_MXFP8, MEGA_MOE_NVFP4, MEGA_MOE_BF16) +} + + +def check_deep_gemm_mega_moe_supported( + device: torch.device, hidden_size: int, intermediate_size: int +) -> None: + """Raise unless DeepGEMM MegaMoE can run on ``device`` with these shapes. + + Raises: + NotImplementedError: If the device architecture is unsupported. + ValueError: If the shapes are unsupported. + + """ + if not current_platform.is_device_capability_family( + 100, device_id=device.index or 0 + ): + raise NotImplementedError("DeepGEMM MegaMoE requires SM100 GPUs.") + if hidden_size % 128 != 0 or intermediate_size % 128 != 0: + raise ValueError( + "DeepGEMM MegaMoE requires hidden and intermediate sizes " + "to be multiples of 128." + ) + + +def _linear_weight_scale(linear: nn.Module) -> torch.Tensor | None: + scale = getattr(linear, "weight_scale", None) + return getattr(linear, "weight_scale_inv", None) if scale is None else scale + + +def _matches_weight_spec(linear: nn.Module, spec: MegaMoEWeightSpec) -> bool: + """Whether a checkpoint linear can be written into a ``spec`` weight.""" + weight, scale = linear.weight, _linear_weight_scale(linear) + if weight.dtype != spec.dtype or weight.dim() != 2: + return False + if spec.scale_group_k is None: + return scale is None + if scale is None or scale.dim() != 2: + return False + (rows, cols), (scale_rows, scale_cols) = weight.shape, scale.shape + scale_dtypes = {spec.scale_dtype} + if spec.scale_dtype == torch.float8_e8m0fnu: + scale_dtypes.add(torch.uint8) + return ( + scale.dtype in scale_dtypes + and rows % scale_rows == 0 + and cols % scale_cols == 0 + and (cols // scale_cols) % spec.scale_group_k == 0 + ) + + +class DeepGemmMegaMoEExperts(nn.Module): + """Routed experts, and optionally the shared expert, as one MegaMoE launch. + + Parameters are registered in the kernel layout and every checkpoint shard is + written straight into it, so loading (and reloading) needs no post-load + transform. Whether the shared expert runs inside the launch is decided here, + before any weight loads; a fused layer owns the shared-expert weights and + the caller drops its shared MLP. + + NVFP4 experts also take per-expert ``weight_scale_2`` (``w1`` and ``w3`` + must match: one L1 scale per expert) and ``w2`` ``input_scale``; the kernel + quantizes every L2 input with the max of the latter over all loaded + experts. Dispatch uses per-token scales, so ``w1``/``w3`` ``input_scale`` + are not loaded. + + Args: + num_experts: Physical experts across EP ranks, redundant included. + activation: Activation fused into the L1 epilogue. + weight_format: Routed-expert precision, a row of the feature matrix. + shared_experts: Shared MLP to fuse if it is unsharded and in one of + the ``weight_format.shared`` precisions; ``None`` never fuses. + + """ + + _symm_buffers: ClassVar[dict[tuple, Any]] = {} + _synchronized_ep_groups: ClassVar[set[tuple[int, int]]] = set() + + def __init__( self, - layer: nn.Module, + vllm_config: VllmConfig, *, + num_experts: int, num_local_experts: int, + experts_start_idx: int, + top_k: int, hidden_size: int, intermediate_size: int, - 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 + activation: MegaMoEActivation, + weight_format: MegaMoEWeightFormat, + shared_experts: nn.Module | None = None, + num_logical_experts: int | None = None, + fast_math: bool = True, + prefix: str = "", + ): + super().__init__() + if activation.name not in weight_format.activations: + raise NotImplementedError( + f"DeepGEMM MegaMoE {weight_format.name} experts support " + f"{sorted(weight_format.activations)}, not {activation.name!r}." + ) + 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() + ) + check_deep_gemm_mega_moe_supported(device, hidden_size, intermediate_size) - def load_weight( - self, - 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 + self.prefix = prefix + self.capture_fn: Callable[[torch.Tensor], None] | None = None + self.num_experts = num_experts + self.num_local_experts = num_local_experts + self.experts_start_idx = experts_start_idx + self.experts_end_idx = experts_start_idx + num_local_experts + self.num_logical_experts = ( + num_logical_experts if num_logical_experts is not None else num_experts + ) + self.top_k = top_k + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens + self.weight_format = weight_format + self.activation = activation + self.fast_math = fast_math + self._kernel_kwargs = activation.kernel_kwargs() + self.eplb_state = EplbLayerState() + # L2 input scale of every logical expert; the NVFP4 kernel takes the max. + self._w2_input_scales = torch.zeros(self.num_logical_experts, device="cpu") + self.l2_activation_scale = 1.0 + # The first of each expert's w1/w3 weight_scale_2 until the other loads. + self._pending_l1_scales: dict[int, tuple[str, float]] = {} - def supports_shared_experts(self, shared_experts: nn.Module) -> bool: - """Whether ``shared_experts`` can be fused into the MegaMoE kernel. + self.num_shared_experts, self.shared_spec = self._fusable_shared_experts( + shared_experts + ) + self._create_weights() + 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}) + if not shared and name.endswith("weight_scale"): + param.quant_method = "block" - Decided at construction from the installed kernel API and the shared - MLP's parameters, so the fused layout is fixed before any weight loads. + # 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 + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + @property + def has_fused_shared_experts(self) -> bool: + return self.num_shared_experts > 0 + + def _fusable_shared_experts( + self, shared_experts: nn.Module | None + ) -> tuple[int, MegaMoEWeightSpec | None]: + """How many shared experts the launch fuses, and in which precision.""" + if shared_experts is None: + return 0, None + gate_up, down = shared_experts.gate_up_proj, shared_experts.down_proj + hidden, inter = self.hidden_size, self.intermediate_size + shared_size = down.weight.shape[-1] + spec = next( + ( + spec + for spec in self.weight_format.shared + if all(_matches_weight_spec(lin, spec) for lin in (gate_up, down)) + ), + None, + ) + if ( + tuple(gate_up.weight.shape) != (2 * shared_size, hidden) + or tuple(down.weight.shape) != (hidden, shared_size) + or shared_size % inter + ): + reason = ( + f"needs an unsharded shared MLP over hidden size {hidden} with a " + f"multiple of {inter} intermediate; got {tuple(gate_up.weight.shape)}" + f" and {tuple(down.weight.shape)}" + ) + elif spec is None: + scale = _linear_weight_scale(gate_up) + supported = ", ".join( + f"{s.dtype}{' with 1x32 UE8M0 scales' if s.scale_group_k else ''}" + for s in self.weight_format.shared + ) + reason = ( + f"{self.weight_format.name} experts fuse {supported} shared " + f"experts; got weight {gate_up.weight.dtype} and scale " + f"{None if scale is None else (tuple(scale.shape), scale.dtype)}" + ) + elif not self._kernel_fuses_shared_experts(spec): + reason = ( + "the installed DeepGEMM Python API is older than the vLLM source. " + "Rebuild the vendored _deep_gemm_C extension to enable it" + ) + else: + return shared_size // inter, spec + logger.warning_once("Disabling MegaMoE shared-expert fusion: %s.", reason) + return 0, None + + def _kernel_fuses_shared_experts(self, spec: MegaMoEWeightSpec) -> bool: + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + fmt = self.weight_format + buffer_api = "get_symm_buffer_for_mega_moe" + block_m_api = "get_block_m_for_mega_moe" + if fmt.uses_nvfp4_api: + buffer_api = "NVFP4SymmBuffer" + block_m_api = "get_block_m_for_nvfp4_mega_moe" + try: + api = { + *signature(getattr(deep_gemm, buffer_api)).parameters, + *signature(getattr(deep_gemm, fmt.kernel)).parameters, + } + except (AttributeError, TypeError, ValueError): + return False + # Scaled shared inputs are laid out by the runtime BLOCK_M. + return { + "num_shared_experts", + "shared_l1_weights", + "shared_l2_weights", + } <= api and (spec.scale_group_k is None or hasattr(deep_gemm, block_m_api)) + + def _create_weights(self) -> None: + """Register the routed and fused shared-expert parameters. + + Routed parameters are indexed by local expert on dim 0, which EPLB moves. """ - return False + e, h, i = self.num_local_experts, self.hidden_size, self.intermediate_size + fmt = self.weight_format + self._register_weight("w13", (e, 2 * i, h), fmt.routed) + self._register_weight("w2", (e, h, i), fmt.routed) + if fmt.routed.has_global_scale: + # w13_weight_scale_2 is the kernel's L1 alpha (per-token dispatch). + for name in ("w13_weight_scale_2", "w2_weight_scale_2", "w2_input_scale"): + scale = torch.zeros(e, dtype=torch.float32) + self.register_parameter(name, nn.Parameter(scale, requires_grad=False)) + # w2 global weight scale times the L2 activation scale, per launch. + l2_alpha = torch.zeros(e, dtype=torch.float32) + self.register_buffer("_l2_alpha", l2_alpha, persistent=False) + if self.shared_spec is not None: + # The shared expert has the layout of one routed expert slice. + s = i * self.num_shared_experts + self._register_weight("shared_l1", (2 * s, h), self.shared_spec) + self._register_weight("shared_l2", (h, s), self.shared_spec) - def run_mega_moe( - self, - *, - y: torch.Tensor, - l1_weights: tuple[torch.Tensor, torch.Tensor], - l2_weights: tuple[torch.Tensor, torch.Tensor], - symm_buffer, - activation_clamp: float | None, - fast_math: bool, - activation: str | None = None, - activation_alpha: float | None = None, - activation_beta: float | None = None, - shared_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None, - shared_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None, - ) -> None: - raise NotImplementedError - - def supports_bf16_mega_gate(self) -> bool: - """Whether this backend can run the DeepGEMM bf16 MegaMoE gate kernel. - - Defaults to ``False`` (the symbol is not guaranteed to exist in every - deep_gemm build). Backends that target a deep_gemm build exposing the - ``bf16_mega_gate`` symbol override this and :meth:`bf16_mega_gate`. - """ - return False - - def bf16_mega_gate( - self, - *, - x: torch.Tensor, - weight: torch.Tensor, - num_topk: int, - scoring_func: str, - routed_scaling_factor: float, - ep_rank: int, - bias: torch.Tensor | None = None, - image_bias: torch.Tensor | None = None, - image_token_mask: torch.Tensor | None = None, - fix_routing_mask: torch.Tensor | None = None, - unmapped_topk_idx: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - raise NotImplementedError - - -class DeepGemmSm100MegaMoEBackend(DeepGemmMegaMoEBackend): - """Default MegaMoE backend targeting SM100 via deep_gemm fp8_fp4_mega_moe.""" - - mma_type = "fp8xfp4" - hidden_quant = kMxfp8Dynamic - - def create_weights( - self, - layer: nn.Module, - *, - num_local_experts: int, - hidden_size: int, - intermediate_size: int, - num_shared_experts: int, + def _register_weight( + self, name: str, shape: tuple[int, ...], spec: MegaMoEWeightSpec ) -> None: def register(name: str, data: torch.Tensor) -> None: - layer.register_parameter(name, nn.Parameter(data, requires_grad=False)) + self.register_parameter(name, nn.Parameter(data, requires_grad=False)) - # 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)) + *outer, n, k = shape register( - "w13_weight_scale", - torch.zeros(e, h // 128, 2 * i, dtype=torch.int32).transpose(1, 2), - ) - 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), - ) - 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), + f"{name}_weight", + torch.zeros(*outer, n, k // spec.values_per_element, dtype=spec.dtype), ) + if spec.scale_group_k is not None: + # Block scales, MN-major, four per int32 along K. + scale_k = k // (4 * spec.scale_group_k) + scale = torch.zeros(*outer, scale_k, n, dtype=torch.int32) + register(f"{name}_weight_scale", scale.transpose(-1, -2)) - def load_weight( + def _kernel_weights( + self, name: str + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + weight = getattr(self, f"{name}_weight") + scale = getattr(self, f"{name}_weight_scale", None) + return weight if scale is None else (weight, scale) + + @property + def l1_weights(self) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + return self._kernel_weights("w13") + + @property + def l2_weights(self) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + return self._kernel_weights("w2") + + def _map_global_expert_id(self, expert_id: int) -> list[int]: + """Local slots that logical expert ``expert_id`` occupies on this rank.""" + physical_ids: list[int] = [] + for p in range(self.experts_start_idx, self.experts_end_idx): + if p % self.num_logical_experts == expert_id: + physical_ids.append(p - self.experts_start_idx) + return physical_ids + + def weight_loader( + self, + param: nn.Parameter, + loaded_weight: torch.Tensor, + weight_name: str, + shard_id: str, + expert_id: int, + return_success: bool = False, + ) -> bool | None: + if not loaded_weight.is_meta: + self._record_global_scale(weight_name, loaded_weight, shard_id, expert_id) + local_expert_ids = self._map_global_expert_id(expert_id) + 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 + is_global_scale = weight_name.endswith(("weight_scale_2", "input_scale")) + is_scale = weight_name.endswith("weight_scale") + for local_expert_id in local_expert_ids: + if is_global_scale: + param.data[local_expert_id].copy_(loaded_weight.reshape(())) + else: + self._load_weight( + param.data[local_expert_id], loaded_weight, shard_id, is_scale + ) + return True if return_success else None + + def _record_global_scale( + self, weight_name: str, scale: torch.Tensor, shard_id: str, expert_id: int + ) -> None: + """Track NVFP4 global scales the kernel takes one of per launch or expert. + + Every rank sees the scales of all experts, so ranks agree and EPLB can + move experts. + """ + if weight_name.endswith("w2_input_scale"): + self._w2_input_scales[expert_id] = scale.max() + self.l2_activation_scale = float(self._w2_input_scales.max()) or 1.0 + elif weight_name.endswith("w13_weight_scale_2"): + # Compare w1 and w3 of the same load, not the live (maybe stale) slot. + prev = self._pending_l1_scales.pop(expert_id, None) + value = float(scale) + if prev is None or prev[0] == shard_id: + self._pending_l1_scales[expert_id] = (shard_id, value) + elif not math.isclose(value, prev[1], rel_tol=1e-5): + raise ValueError( + f"MegaMoE NVFP4 expert {expert_id}: {prev[0]} weight_scale_2 " + f"{prev[1]} != {shard_id} weight_scale_2 {value}; the kernel " + "takes one L1 scale per expert." + ) + + 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}" + + 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._load_weight(param.data, loaded_weight, shard_id, is_scale) + + def _load_weight( self, dst: torch.Tensor, src: torch.Tensor, shard_id: str, is_scale: bool, ) -> None: + """Write checkpoint shard ``src`` into ``dst`` in the kernel layout. + + ``dst`` is one routed expert's slice or a shared-expert parameter. + """ 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) @@ -182,7 +571,7 @@ class DeepGemmSm100MegaMoEBackend(DeepGemmMegaMoEBackend): @staticmethod 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.""" + """Expand block scales to the kernel granularity, four per int32 along K.""" sf = sf.view(torch.uint8) if rows % sf.shape[0] or cols % sf.shape[1]: raise ValueError( @@ -214,163 +603,255 @@ class DeepGemmSm100MegaMoEBackend(DeepGemmMegaMoEBackend): 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 - + def get_symm_buffer(self): 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.", + group = get_ep_group().device_group + key = ( + id(group), + torch.accelerator.current_device_index(), + self.num_experts, + self.max_num_tokens, + self.top_k, + self.hidden_size, + self.intermediate_size, + self.num_shared_experts, + None if self.shared_spec is None else self.shared_spec.dtype, + self.weight_format.mma_type, + self.activation.name, + ) + symm_buffer = self._symm_buffers.get(key) + if symm_buffer is None: + args = ( + group, + self.num_experts, + self.max_num_tokens, + self.top_k, + self.hidden_size, + self.intermediate_size, ) - return False - return True + if self.weight_format.uses_nvfp4_api: + symm_buffer = deep_gemm.NVFP4SymmBuffer( + *args, + num_shared_experts=self.num_shared_experts, + shared_dtype=torch.float8_e4m3fn + if self.shared_spec is None + else self.shared_spec.dtype, + ) + else: + symm_buffer = deep_gemm.get_symm_buffer_for_mega_moe( + *args, + num_shared_experts=self.num_shared_experts, + mma_type=self.weight_format.mma_type, + activation=self.activation.name, + ) + self._symm_buffers[key] = symm_buffer + return symm_buffer - @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 + def get_block_m(self, num_tokens: int) -> int: + """BLOCK_M the MegaMoE scheduler runs ``num_tokens`` local tokens with.""" + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + args = ( + get_ep_group().world_size, + self.num_experts, + self.get_symm_buffer().num_max_tokens_per_rank, + num_tokens, + self.top_k, + ) + if self.weight_format.uses_nvfp4_api: + return deep_gemm.get_block_m_for_nvfp4_mega_moe(*args) + return deep_gemm.get_block_m_for_mega_moe(*args, self.weight_format.mma_type) + + def set_eplb_state( + self, + moe_layer_idx: int, + expert_load_view: torch.Tensor, + logical_to_physical_map: torch.Tensor, + logical_replica_count: torch.Tensor, + ) -> None: + self.eplb_state.set_layer_state( + moe_layer_idx, + expert_load_view, + logical_to_physical_map, + logical_replica_count, ) - def run_mega_moe( + def get_expert_weights(self) -> list[torch.Tensor]: + 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 + if t.is_contiguous(): + return t.view(self.num_local_experts, -1) + elif t.dim() == 3 and t.stride(1) == 1 and t.stride(2) == t.shape[1]: + # scales have shape (E, M, N) with memory layout (E, N, M) + back = torch.transpose(t, 1, 2) + assert back.is_contiguous() + return back.view(self.num_local_experts, -1) + + raise AssertionError( + f"MegaMoE EPLB {name}: non-contiguous expert tensor with " + f"unexpected layout shape={tuple(t.shape)} " + f"stride={tuple(t.stride())} dtype={t.dtype}" + ) + + return [ + _to_eplb_view(name, param) + for name, param in self.named_parameters(recurse=False) + if not name.startswith("shared_") + ] + + def update_expert_map(self) -> None: + pass + + @property + def layer_id(self) -> int: + return extract_layer_index(self.prefix) + + def _synchronize_first_launch(self) -> None: + ep_group = get_ep_group() + key = (id(ep_group.cpu_group), torch.accelerator.current_device_index()) + if key in self._synchronized_ep_groups: + return + torch.accelerator.synchronize() + torch.distributed.barrier(group=ep_group.cpu_group) + self._synchronized_ep_groups.add(key) + + def forward( self, - *, - y: torch.Tensor, - l1_weights: tuple[torch.Tensor, torch.Tensor], - l2_weights: tuple[torch.Tensor, torch.Tensor], - symm_buffer, - activation_clamp: float | None, - fast_math: bool, - activation: str | None = None, - activation_alpha: float | None = None, - activation_beta: float | None = None, - shared_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None, - shared_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None, - ) -> None: + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts: nn.Module | None = None, + ) -> torch.Tensor: + """Run the routed and shared experts on this rank's tokens. + + Args: + hidden_states: BF16 ``[num_tokens, hidden_size]`` activations. + topk_weights: ``[num_tokens, top_k]`` routing weights. + topk_ids: ``[num_tokens, top_k]`` logical expert ids. + shared_experts: Unfused shared MLP, run on ``hidden_states`` and + added to the output. Must be ``None`` when this layer fused the + shared expert at construction. + + Returns: + BF16 ``[num_tokens, hidden_size]`` routed plus shared expert output. + + """ + if shared_experts is not None and self.has_fused_shared_experts: + raise ValueError( + f"{self.prefix} fuses the shared expert into MegaMoE; " + "pass shared_experts=None." + ) + num_tokens = hidden_states.shape[0] + if num_tokens > self.max_num_tokens: + raise ValueError( + f"MegaMoE got {num_tokens} tokens, but the symmetric buffer was " + f"sized for {self.max_num_tokens}." + ) + self._synchronize_first_launch() + + is_padding = None + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): + is_padding = get_forward_context().is_padding + if is_padding is not None: + is_padding = is_padding[:num_tokens] + + if self.capture_fn is not None: + self.capture_fn(topk_ids) + + # EPLB: map logical expert IDs to physical replicas and record load. + eplb_state = self.eplb_state + if eplb_state.logical_to_physical_map is not None: + assert eplb_state.expert_load_view is not None + assert eplb_state.logical_replica_count is not None + assert eplb_state.should_record_tensor is not None + if is_padding is not None: + topk_ids = torch.where(is_padding.unsqueeze(1), -1, topk_ids) + topk_ids = eplb_map_to_physical_and_record( + topk_ids=topk_ids, + expert_load_view=eplb_state.expert_load_view, + logical_to_physical_map=eplb_state.logical_to_physical_map, + logical_replica_count=eplb_state.logical_replica_count, + record_enabled=eplb_state.should_record_tensor, + num_unpadded_tokens=eplb_state.num_unpadded_tokens_tensors[ + dbo_current_ubatch_id() + ] + if eplb_state.num_unpadded_tokens_tensors is not None + else None, + ) + from vllm.utils.deep_gemm import _import_deep_gemm deep_gemm = _import_deep_gemm() - kwargs: dict[str, Any] = { - name: value - for name, value in ( - ("activation", activation), - ("activation_alpha", activation_alpha), - ("activation_beta", activation_beta), + fmt = self.weight_format + symm_buffer = self.get_symm_buffer() + shared_x_sf = None + shared_block_m = None + if self.shared_spec is not None and self.shared_spec.scale_group_k: + shared_x_sf = symm_buffer.shared_l1_acts_sf + shared_block_m = self.get_block_m(num_tokens) + kwargs = dict(self._kernel_kwargs) + if fmt.uses_nvfp4_api: + l2_activation_scale = self.l2_activation_scale + prepare_megamoe_nvfp4_inputs( + hidden_states, + topk_weights, + topk_ids, + symm_buffer.x[:num_tokens], + symm_buffer.x_sf[:num_tokens], + symm_buffer.x_scales[:num_tokens], + symm_buffer.topk_idx[:num_tokens], + symm_buffer.topk_weights[:num_tokens], + is_padding=is_padding, + shared_x=symm_buffer.shared_l1_acts[:num_tokens] + if self.has_fused_shared_experts + else None, + shared_x_sf=shared_x_sf, + shared_block_m=shared_block_m, + w2_global_scale=self.w2_weight_scale_2, + l2_alpha=self._l2_alpha, + l2_activation_scale=l2_activation_scale, ) - if value is not None - } - if shared_l1_weights is not None and shared_l2_weights is not None: - kwargs["shared_l1_weights"] = shared_l1_weights - kwargs["shared_l2_weights"] = shared_l2_weights - deep_gemm.fp8_fp4_mega_moe( + kwargs.update( + l1_alpha=self.w13_weight_scale_2, + l2_alpha=self._l2_alpha, + l2_activation_scale=l2_activation_scale, + use_x_scales=True, + ) + else: + quantized = fmt.hidden_quant is not None + prepare_megamoe_inputs( + hidden_states, + topk_weights, + topk_ids, + symm_buffer.x[:num_tokens], + symm_buffer.x_sf[:num_tokens] if quantized else None, + symm_buffer.topk_idx[:num_tokens], + symm_buffer.topk_weights[:num_tokens], + is_padding=is_padding, + shared_x_sf=shared_x_sf, + shared_block_m=shared_block_m, + hidden_quant=fmt.hidden_quant, + ) + + if self.has_fused_shared_experts: + kwargs["shared_l1_weights"] = self._kernel_weights("shared_l1") + kwargs["shared_l2_weights"] = self._kernel_weights("shared_l2") + y = torch.empty_like(hidden_states, dtype=torch.bfloat16) + getattr(deep_gemm, fmt.kernel)( y, - l1_weights, - l2_weights, + self.l1_weights, + self.l2_weights, symm_buffer, - activation_clamp=activation_clamp, - fast_math=fast_math, + fast_math=self.fast_math, **kwargs, ) - - def supports_bf16_mega_gate(self) -> bool: - """SM100 deep_gemm build is expected to provide the bf16 gate kernel.""" - return True - - def bf16_mega_gate( - self, - *, - x: torch.Tensor, - weight: torch.Tensor, - num_topk: int, - scoring_func: str, - routed_scaling_factor: float, - ep_rank: int, - bias: torch.Tensor | None = None, - image_bias: torch.Tensor | None = None, - image_token_mask: torch.Tensor | None = None, - fix_routing_mask: torch.Tensor | None = None, - unmapped_topk_idx: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - from vllm.utils.deep_gemm import bf16_mega_gate - - return bf16_mega_gate( - x, - weight, - num_topk, - scoring_func=scoring_func, - routed_scaling_factor=routed_scaling_factor, - ep_rank=ep_rank, - bias=bias, - image_bias=image_bias, - image_token_mask=image_token_mask, - fix_routing_mask=fix_routing_mask, - unmapped_topk_idx=unmapped_topk_idx, - ) + if shared_experts is not None: + y += shared_experts(hidden_states) + return y -def get_deep_gemm_mega_moe_backend( - device: torch.device, - hidden_size: int, - intermediate_size: int, -) -> DeepGemmMegaMoEBackend: - """Return the DeepGEMM MegaMoE backend for ``device``. - - A successful return implies the backend can run on this device with the - given ``hidden_size`` / ``intermediate_size``. All runtime capability - checks (arch + shape divisibility) live here; there is no separate check - on the returned backend. - - Currently only the SM100 deep_gemm backend is registered here. SM90 and - other backends are expected to be provided by new backends or plugins - that override this factory. - - Raises: - NotImplementedError: If the device architecture is unsupported. - ValueError: If the shapes are unsupported. - - """ - if not current_platform.is_device_capability_family( - 100, device_id=device.index or 0 - ): - raise NotImplementedError("DeepGEMM MegaMoE requires SM100 GPUs.") - if hidden_size % 128 != 0 or intermediate_size % 128 != 0: - raise ValueError( - "DeepGEMM MegaMoE requires hidden and intermediate sizes " - "to be multiples of 128." - ) - return DeepGemmSm100MegaMoEBackend() +DeepGemmMegaMoEExperts.weight_loader.supports_moe_loading = True # type: ignore[attr-defined] diff --git a/vllm/model_executor/layers/fused_moe/prepare_megamoe.py b/vllm/model_executor/layers/fused_moe/prepare_megamoe.py new file mode 100644 index 0000000000..5e8d8cde9a --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/prepare_megamoe.py @@ -0,0 +1,682 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Triton input-staging kernels for DeepGEMM MegaMoE. + +Quantize hidden states to fp8 with E8M0 group scales, to NVFP4 with per-token +global scales, or copy them as BF16, and repack the routing top-k tensors into +the int64/float32 layout that the DeepGEMM MegaMoE kernels consume. +""" + +import torch + +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + FP8_DTYPE, + QuantKey, + kMxfp8Dynamic, +) +from vllm.triton_utils import tl, triton + +_E8M0_SCALE_DTYPES = (torch.uint8, torch.float8_e8m0fnu) +# Inverse of the NVFP4 range: E2M1 max (6) times E4M3 max (448). +_INV_NVFP4_MAX = tl.constexpr(1.0 / 2688.0) + + +@triton.jit +def _quantize_fp8_groups( + hidden, + BLOCK_M: tl.constexpr, + BLOCK_K: tl.constexpr, + GROUP_K: tl.constexpr, + USE_UE8M0: tl.constexpr, +): + """Return fp8 values, fp32 group scales and, for UE8M0, the scale exponents.""" + num_groups: tl.constexpr = BLOCK_K // GROUP_K + hidden_groups = tl.reshape(tl.abs(hidden), [BLOCK_M, num_groups, GROUP_K]) + amax = tl.max(hidden_groups, axis=2) + amax = tl.maximum(amax, 1.0e-4) + + scale = amax / 448.0 + if USE_UE8M0: + scale_bits = scale.to(tl.uint32, bitcast=True) + scale_exp = ((scale_bits >> 23) & 0xFF) + ((scale_bits & 0x7FFFFF) != 0).to( + tl.uint32 + ) + scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254) + quant_scale = (scale_exp << 23).to(tl.float32, bitcast=True) + else: + scale_exp = scale.to(tl.uint32, bitcast=True) + quant_scale = scale + + hidden_groups = tl.reshape(hidden, [BLOCK_M, num_groups, GROUP_K]) + scaled = hidden_groups * (1.0 / quant_scale)[:, :, None] + scaled = tl.reshape(scaled, [BLOCK_M, BLOCK_K]) + return scaled.to(tl.float8e4nv), quant_scale, scale_exp + + +@triton.jit +def _pack_bytes(values, NUM_BYTES: tl.constexpr): + """Pack the last dim of ``[..., NUM_BYTES]`` byte values into int32 words.""" + offsets = tl.arange(0, NUM_BYTES) + return tl.sum(values.to(tl.uint32) << (offsets * 8), axis=-1).to(tl.int32) + + +@triton.jit +def _shared_sf_row(token_id, SHARED_BLOCK_M: tl.constexpr): + """Row of ``token_id`` in the shared-expert MN-major UTCCP scale pages.""" + m_block_id = token_id // SHARED_BLOCK_M + m_in_block = token_id % SHARED_BLOCK_M + aligned_block_m: tl.constexpr = triton.cdiv(SHARED_BLOCK_M, 128) * 128 + transposed_m = ( + (m_in_block // 128) * 128 + (m_in_block % 32) * 4 + (m_in_block % 128) // 32 + ) + return m_block_id * aligned_block_m + transposed_m + + +@triton.jit +def _store_topk( + token_id, + token_mask, + topk_ids, + topk_weights, + is_padding, + topk_idx_out, + topk_weights_out, + topk_ids_stride_m: tl.constexpr, + topk_ids_stride_k: tl.constexpr, + topk_weights_stride_m: tl.constexpr, + topk_weights_stride_k: tl.constexpr, + is_padding_stride_m: tl.constexpr, + topk_idx_stride_m: tl.constexpr, + topk_idx_stride_k: tl.constexpr, + topk_weights_out_stride_m: tl.constexpr, + topk_weights_out_stride_k: tl.constexpr, + top_k: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_TOPK: tl.constexpr, +) -> None: + """Repack top-k routing as int64 ids and fp32 weights; padding is unrouted.""" + topk_offsets = tl.arange(0, BLOCK_TOPK) + topk_mask = token_mask[:, None] & (topk_offsets[None, :] < top_k) + token_is_padding = tl.full((BLOCK_M,), False, tl.int1) + if is_padding is not None: + token_is_padding = tl.load( + is_padding + token_id * is_padding_stride_m, + mask=token_mask, + other=True, + ) + + ids = tl.load( + topk_ids + + token_id[:, None] * topk_ids_stride_m + + topk_offsets[None, :] * topk_ids_stride_k, + mask=topk_mask, + other=0, + ).to(tl.int64) + ids = tl.where(token_is_padding[:, None], -1, ids) + tl.store( + topk_idx_out + + token_id[:, None] * topk_idx_stride_m + + topk_offsets[None, :] * topk_idx_stride_k, + ids, + mask=topk_mask, + ) + + weights = tl.load( + topk_weights + + token_id[:, None] * topk_weights_stride_m + + topk_offsets[None, :] * topk_weights_stride_k, + mask=topk_mask, + other=0.0, + ) + weights = tl.where(token_is_padding[:, None], 0.0, weights) + tl.store( + topk_weights_out + + token_id[:, None] * topk_weights_out_stride_m + + topk_offsets[None, :] * topk_weights_out_stride_k, + weights, + mask=topk_mask, + ) + + +@triton.jit(do_not_specialize=["num_tokens"]) +def _prepare_megamoe_inputs_kernel( + hidden_states, + x_out, + x_sf, + shared_x_sf, + topk_ids, + topk_weights, + is_padding, + topk_idx_out, + topk_weights_out, + hidden_stride_m: tl.constexpr, + hidden_stride_k: tl.constexpr, + x_stride_m: tl.constexpr, + x_stride_k: tl.constexpr, + x_sf_stride_m: tl.constexpr, + x_sf_stride_k: tl.constexpr, + shared_x_sf_stride_m: tl.constexpr, + shared_x_sf_stride_k: tl.constexpr, + topk_ids_stride_m: tl.constexpr, + topk_ids_stride_k: tl.constexpr, + topk_weights_stride_m: tl.constexpr, + topk_weights_stride_k: tl.constexpr, + is_padding_stride_m: tl.constexpr, + topk_idx_stride_m: tl.constexpr, + topk_idx_stride_k: tl.constexpr, + topk_weights_out_stride_m: tl.constexpr, + topk_weights_out_stride_k: tl.constexpr, + num_tokens, + hidden_size: tl.constexpr, + top_k: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_K: tl.constexpr, + GROUP_K: tl.constexpr, + BLOCK_TOPK: tl.constexpr, + SHARED_BLOCK_M: tl.constexpr, + QUANTIZE: tl.constexpr, + USE_UE8M0: tl.constexpr, +) -> None: + token_id = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) + token_mask = token_id < num_tokens + k_block_id = tl.program_id(1) + + k_offsets = k_block_id * BLOCK_K + tl.arange(0, BLOCK_K) + k_mask = token_mask[:, None] & (k_offsets[None, :] < hidden_size) + hidden = tl.load( + hidden_states + + token_id[:, None] * hidden_stride_m + + k_offsets[None, :] * hidden_stride_k, + mask=k_mask, + other=0.0, + ).to(tl.float32) + + if not QUANTIZE: + tl.store( + x_out + token_id[:, None] * x_stride_m + k_offsets[None, :] * x_stride_k, + hidden.to(x_out.dtype.element_ty), + mask=k_mask, + ) + else: + num_groups: tl.constexpr = BLOCK_K // GROUP_K + fp8, quant_scale, scale_exp = _quantize_fp8_groups( + hidden, BLOCK_M, BLOCK_K, GROUP_K, USE_UE8M0 + ) + tl.store( + x_out + token_id[:, None] * x_stride_m + k_offsets[None, :] * x_stride_k, + fp8, + mask=k_mask, + ) + + if USE_UE8M0: + packed_scale = _pack_bytes(scale_exp, num_groups) + tl.store( + x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k, + packed_scale, + mask=token_mask, + ) + else: + sf_offset = tl.arange(0, num_groups) + tl.store( + x_sf + + token_id[:, None] * x_sf_stride_m + + k_block_id * x_sf_stride_k + + sf_offset[None, :], + quant_scale, + mask=token_mask[:, None], + ) + + # DeepGEMM's SM100 shared-expert TMA loads require the activation scales + # in an MN-major layout whose row permutation depends on the MegaMoE + # scheduler's runtime BLOCK_M. Write that view while the packed UE8M0 scale + # is already resident, avoiding another kernel and temporary tensor. + if USE_UE8M0 and shared_x_sf is not None: + shared_row = _shared_sf_row(token_id, SHARED_BLOCK_M) + tl.store( + shared_x_sf + + shared_row * shared_x_sf_stride_m + + k_block_id * shared_x_sf_stride_k, + packed_scale, + mask=token_mask, + ) + + if k_block_id == 0: + _store_topk( + token_id, + token_mask, + topk_ids, + topk_weights, + is_padding, + topk_idx_out, + topk_weights_out, + topk_ids_stride_m, + topk_ids_stride_k, + topk_weights_stride_m, + topk_weights_stride_k, + is_padding_stride_m, + topk_idx_stride_m, + topk_idx_stride_k, + topk_weights_out_stride_m, + topk_weights_out_stride_k, + top_k, + BLOCK_M, + BLOCK_TOPK, + ) + + +@triton.jit(do_not_specialize=["num_tokens", "l2_activation_scale"]) +def _prepare_megamoe_nvfp4_inputs_kernel( + hidden_states, + x_out, + x_sf, + x_scales, + shared_x, + shared_x_sf, + topk_ids, + topk_weights, + is_padding, + topk_idx_out, + topk_weights_out, + w2_global_scale, + l2_alpha, + hidden_stride_m: tl.constexpr, + hidden_stride_k: tl.constexpr, + x_stride_m: tl.constexpr, + x_sf_stride_m: tl.constexpr, + shared_x_stride_m: tl.constexpr, + shared_x_sf_stride_m: tl.constexpr, + shared_x_sf_stride_k: tl.constexpr, + topk_ids_stride_m: tl.constexpr, + topk_ids_stride_k: tl.constexpr, + topk_weights_stride_m: tl.constexpr, + topk_weights_stride_k: tl.constexpr, + is_padding_stride_m: tl.constexpr, + topk_idx_stride_m: tl.constexpr, + topk_idx_stride_k: tl.constexpr, + topk_weights_out_stride_m: tl.constexpr, + topk_weights_out_stride_k: tl.constexpr, + num_tokens, + l2_activation_scale, + hidden_size: tl.constexpr, + top_k: tl.constexpr, + num_local_experts: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_TOPK: tl.constexpr, + BLOCK_E: tl.constexpr, + SHARED_BLOCK_M: tl.constexpr, + SHARED_MXFP8: tl.constexpr, +) -> None: + if l2_alpha is not None: + experts = tl.arange(0, BLOCK_E) + expert_mask = (experts < num_local_experts) & (tl.program_id(0) == 0) + alpha = tl.load(w2_global_scale + experts, mask=expert_mask) + tl.store(l2_alpha + experts, alpha * l2_activation_scale, mask=expert_mask) + + token_id = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) + token_mask = token_id < num_tokens + k_offsets = tl.arange(0, BLOCK_K) + hidden_ptrs = hidden_states + token_id[:, None] * hidden_stride_m + + amax = tl.zeros((BLOCK_M,), tl.float32) + for k_block_id in range(hidden_size // BLOCK_K): + hidden = tl.load( + hidden_ptrs + (k_block_id * BLOCK_K + k_offsets[None, :]) * hidden_stride_k, + mask=token_mask[:, None], + other=0.0, + ).to(tl.float32) + amax = tl.maximum(amax, tl.max(tl.abs(hidden), axis=1)) + x_scale = tl.maximum(amax, 1.0e-12) * _INV_NVFP4_MAX + tl.store(x_scales + token_id, x_scale, mask=token_mask) + + for k_block_id in range(hidden_size // BLOCK_K): + k_ids = k_block_id * BLOCK_K + k_offsets + hidden = tl.load( + hidden_ptrs + k_ids[None, :] * hidden_stride_k, + mask=token_mask[:, None], + other=0.0, + ) + + # E2M1 values with one E4M3 scale per 16, relative to the token scale. + values = tl.reshape( + tl.div_rn(hidden.to(tl.float32), x_scale[:, None]), + [BLOCK_M, BLOCK_K // 16, 16], + ) + block_scale = tl.max(tl.abs(values), axis=2) * (1.0 / 6.0) + block_scale = tl.minimum(tl.maximum(block_scale, 2.0**-9), 448.0) + block_scale = block_scale.to(tl.float8e4nv) + values = tl.div_rn(values, block_scale.to(tl.float32)[:, :, None]) + magnitude = tl.abs(values) + # Round to nearest even on the grid {0, 0.5, 1, 1.5, 2, 3, 4, 6}. + code = ( + (magnitude > 0.25).to(tl.int32) + + (magnitude >= 0.75).to(tl.int32) + + (magnitude > 1.25).to(tl.int32) + + (magnitude >= 1.75).to(tl.int32) + + (magnitude > 2.5).to(tl.int32) + + (magnitude >= 3.5).to(tl.int32) + + (magnitude > 5.0).to(tl.int32) + ) + code = tl.where((values < 0) & (code != 0), code | 8, code) + low, high = tl.split(tl.reshape(code, [BLOCK_M, BLOCK_K // 2, 2])) + packed_offsets = k_block_id * (BLOCK_K // 2) + tl.arange(0, BLOCK_K // 2) + tl.store( + x_out + token_id[:, None] * x_stride_m + packed_offsets[None, :], + (low | (high << 4)).to(tl.uint8), + mask=token_mask[:, None], + ) + sf_offsets = k_block_id * (BLOCK_K // 64) + tl.arange(0, BLOCK_K // 64) + sf_bytes = tl.reshape( + block_scale.to(tl.uint8, bitcast=True), [BLOCK_M, BLOCK_K // 64, 4] + ) + tl.store( + x_sf + token_id[:, None] * x_sf_stride_m + sf_offsets[None, :], + _pack_bytes(sf_bytes, 4), + mask=token_mask[:, None], + ) + + if shared_x is not None: + shared_ptrs = shared_x + token_id[:, None] * shared_x_stride_m + if SHARED_MXFP8: + fp8, _, scale_exp = _quantize_fp8_groups( + hidden.to(tl.float32), BLOCK_M, BLOCK_K, 32, True + ) + tl.store(shared_ptrs + k_ids[None, :], fp8, mask=token_mask[:, None]) + tl.store( + shared_x_sf + + _shared_sf_row(token_id, SHARED_BLOCK_M) * shared_x_sf_stride_m + + k_block_id * shared_x_sf_stride_k, + _pack_bytes(scale_exp, BLOCK_K // 32), + mask=token_mask, + ) + else: + tl.store(shared_ptrs + k_ids[None, :], hidden, mask=token_mask[:, None]) + + _store_topk( + token_id, + token_mask, + topk_ids, + topk_weights, + is_padding, + topk_idx_out, + topk_weights_out, + topk_ids_stride_m, + topk_ids_stride_k, + topk_weights_stride_m, + topk_weights_stride_k, + is_padding_stride_m, + topk_idx_stride_m, + topk_idx_stride_k, + topk_weights_out_stride_m, + topk_weights_out_stride_k, + top_k, + BLOCK_M, + BLOCK_TOPK, + ) + + +def _resolve_hidden_quant(hidden_quant: QuantKey, block_k: int) -> tuple[int, bool]: + """Map a hidden-state ``QuantKey`` to the kernel's ``(GROUP_K, USE_UE8M0)``.""" + scale = hidden_quant.scale + group_shape = scale.group_shape + if ( + hidden_quant.dtype != FP8_DTYPE + or not hidden_quant.symmetric + or hidden_quant.scale2 is not None + or scale.static + or not group_shape.is_per_group() + or block_k % group_shape.col != 0 + or scale.dtype not in (*_E8M0_SCALE_DTYPES, torch.float32) + ): + raise ValueError( + "MegaMoE input staging requires symmetric dynamic fp8 " + f"per-group quantization with a group size dividing {block_k} and " + f"E8M0 or fp32 scales, got {hidden_quant}." + ) + return group_shape.col, scale.dtype in _E8M0_SCALE_DTYPES + + +def prepare_megamoe_inputs( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + x: torch.Tensor, + x_sf: torch.Tensor | None, + topk_idx_out: torch.Tensor, + topk_weights_out: torch.Tensor, + is_padding: torch.Tensor | None = None, + shared_x_sf: torch.Tensor | None = None, + shared_block_m: int | None = None, + hidden_quant: QuantKey | None = kMxfp8Dynamic, +) -> None: + """Quantize hidden states and repack top-k routing for DeepGEMM MegaMoE. + + Args: + hidden_states: Input activations of shape ``[num_tokens, hidden_dim]``. + topk_weights: Router top-k weights of shape ``[num_tokens, top_k]``. + topk_ids: Router top-k expert ids of shape ``[num_tokens, top_k]``. + x: Output buffer for the dispatched hidden states. + x_sf: Output buffer for the hidden-state scale factors; ``None`` when + ``hidden_quant`` is ``None``. + topk_idx_out: Output buffer for the repacked top-k expert ids. + topk_weights_out: Output buffer for the repacked top-k weights. + is_padding: Optional per-token mask; padded tokens are not routed. + shared_x_sf: Optional output buffer for shared-expert scale factors. + Must be given together with ``shared_block_m``. + shared_block_m: Block M used to lay out ``shared_x_sf``. + hidden_quant: Hidden-state quantization scheme. E8M0 scales are packed + four per int32 into ``x_sf``; fp32 scales are stored directly. + ``None`` copies the hidden states into ``x`` unquantized. + + """ + num_tokens, hidden_size = hidden_states.shape + if num_tokens == 0: + return + if hidden_size % 128 != 0: + raise ValueError( + "MegaMoE input staging requires hidden_size to be a multiple of 128." + ) + block_k = 128 + quantize = hidden_quant is not None + if quantize != (x_sf is not None): + raise ValueError("MegaMoE input staging needs x_sf iff hidden_quant is set.") + group_k, use_ue8m0 = ( + (block_k, False) + if hidden_quant is None + else _resolve_hidden_quant(hidden_quant, block_k) + ) + top_k = topk_ids.shape[1] + if topk_weights.shape != topk_ids.shape: + raise ValueError( + "MegaMoE input staging requires topk_weights and " + "topk_ids to have the same shape." + ) + if (shared_x_sf is None) != (shared_block_m is None): + raise ValueError( + "MegaMoE shared input staging requires both shared_x_sf and shared_block_m." + ) + if shared_x_sf is not None and not use_ue8m0: + raise ValueError( + "MegaMoE shared input staging currently requires " + "UE8M0-packed hidden scales." + ) + if shared_x_sf is not None: + assert shared_block_m is not None + if shared_block_m <= 0: + raise ValueError("MegaMoE shared_block_m must be positive.") + expected_sf_k = hidden_size // 128 + if shared_x_sf.ndim != 2 or shared_x_sf.shape[1] != expected_sf_k: + raise ValueError( + "MegaMoE shared_x_sf must have shape " + f"(*, {expected_sf_k}), got {tuple(shared_x_sf.shape)}." + ) + aligned_block_m = triton.cdiv(shared_block_m, 128) * 128 + required_rows = triton.cdiv(num_tokens, shared_block_m) * aligned_block_m + if shared_x_sf.shape[0] < required_rows: + raise ValueError( + "MegaMoE shared_x_sf has insufficient rows: requires " + f"{required_rows}, got {shared_x_sf.shape[0]}." + ) + + # On GB200, eight-row tiles win from 64 tokens; keep smaller batches untiled. + block_m = 8 if num_tokens >= 64 else 1 + grid = (triton.cdiv(num_tokens, block_m), triton.cdiv(hidden_size, block_k)) + block_topk = triton.next_power_of_2(top_k) + padding_stride_m = is_padding.stride(0) if is_padding is not None else 0 + _prepare_megamoe_inputs_kernel[grid]( + hidden_states, + x, + x_sf, + shared_x_sf, + topk_ids, + topk_weights, + is_padding, + topk_idx_out, + topk_weights_out, + hidden_states.stride(0), + hidden_states.stride(1), + x.stride(0), + x.stride(1), + x_sf.stride(0) if x_sf is not None else 0, + x_sf.stride(1) if x_sf is not None else 0, + shared_x_sf.stride(0) if shared_x_sf is not None else 0, + shared_x_sf.stride(1) if shared_x_sf is not None else 0, + topk_ids.stride(0), + topk_ids.stride(1), + topk_weights.stride(0), + topk_weights.stride(1), + padding_stride_m, + topk_idx_out.stride(0), + topk_idx_out.stride(1), + topk_weights_out.stride(0), + topk_weights_out.stride(1), + num_tokens, + hidden_size, + top_k, + BLOCK_M=block_m, + BLOCK_K=block_k, + GROUP_K=group_k, + BLOCK_TOPK=block_topk, + SHARED_BLOCK_M=shared_block_m or 1, + QUANTIZE=quantize, + USE_UE8M0=use_ue8m0, + num_warps=4, + ) + + +def prepare_megamoe_nvfp4_inputs( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + x: torch.Tensor, + x_sf: torch.Tensor, + x_scales: torch.Tensor, + topk_idx_out: torch.Tensor, + topk_weights_out: torch.Tensor, + is_padding: torch.Tensor | None = None, + shared_x: torch.Tensor | None = None, + shared_x_sf: torch.Tensor | None = None, + shared_block_m: int | None = None, + w2_global_scale: torch.Tensor | None = None, + l2_alpha: torch.Tensor | None = None, + l2_activation_scale: float = 1.0, +) -> None: + """Stage DeepGEMM NVFP4 MegaMoE inputs (``kNvfp4DynamicToken``). + + Each token is scaled by its own fp32 ``x_scales`` (row amax / 2688) and + quantized to E2M1 with one E4M3 scale per 16 values. + + Args: + hidden_states: BF16 activations of shape ``[num_tokens, hidden_dim]``. + topk_weights: Router top-k weights of shape ``[num_tokens, top_k]``. + topk_ids: Router top-k expert ids of shape ``[num_tokens, top_k]``. + x: Packed E2M1 output, ``[num_tokens, hidden_dim // 2]``. + x_sf: E4M3 scales packed four per int32, ``[num_tokens, hidden_dim // 64]``. + x_scales: Per-token fp32 dequantization scales, ``[num_tokens]``. + topk_idx_out: Output buffer for the repacked top-k expert ids. + topk_weights_out: Output buffer for the repacked top-k weights. + is_padding: Optional per-token mask; padded tokens are not routed. + shared_x: Optional shared-expert input, fp8 (MXFP8 with ``shared_x_sf``) + or a BF16 copy of the hidden states. + shared_x_sf: UE8M0 scales of an fp8 ``shared_x`` in the shared-expert + TMA layout for ``shared_block_m``. + shared_block_m: Block M used to lay out ``shared_x_sf``. + w2_global_scale: Per-local-expert L2 weight global scales. + l2_alpha: Output ``w2_global_scale * l2_activation_scale``, the L2 + dequantization scale of each local expert. Written even without + local tokens, as other ranks' tokens still reach these experts. + l2_activation_scale: Global scale the kernel quantizes L2 inputs with. + + """ + num_tokens, hidden_size = hidden_states.shape + if hidden_size % 128 != 0: + raise ValueError( + "MegaMoE input staging requires hidden_size to be a multiple of 128." + ) + if topk_weights.shape != topk_ids.shape: + raise ValueError( + "MegaMoE input staging requires topk_weights and " + "topk_ids to have the same shape." + ) + shared_mxfp8 = shared_x is not None and shared_x.dtype == FP8_DTYPE + if shared_mxfp8 != (shared_x_sf is not None) or (shared_x_sf is None) != ( + shared_block_m is None + ): + raise ValueError( + "MegaMoE NVFP4 staging needs shared_x_sf and shared_block_m exactly " + "for an fp8 shared_x." + ) + if (w2_global_scale is None) != (l2_alpha is None): + raise ValueError("MegaMoE NVFP4 staging needs w2_global_scale with l2_alpha.") + for name, tensor in (("x", x), ("x_sf", x_sf), ("shared_x", shared_x)): + if tensor is not None and tensor.stride(-1) != 1: + raise ValueError(f"MegaMoE NVFP4 staging requires a row-major {name}.") + + block_m = 8 if num_tokens >= 64 else 1 + top_k = topk_ids.shape[1] + num_local_experts = 0 if l2_alpha is None else l2_alpha.numel() + # One program even without tokens, to write l2_alpha. + grid = (max(triton.cdiv(num_tokens, block_m), 1),) + _prepare_megamoe_nvfp4_inputs_kernel[grid]( + hidden_states, + x.view(torch.uint8), + x_sf, + x_scales, + shared_x, + shared_x_sf, + topk_ids, + topk_weights, + is_padding, + topk_idx_out, + topk_weights_out, + w2_global_scale, + l2_alpha, + hidden_states.stride(0), + hidden_states.stride(1), + x.stride(0), + x_sf.stride(0), + shared_x.stride(0) if shared_x is not None else 0, + shared_x_sf.stride(0) if shared_x_sf is not None else 0, + shared_x_sf.stride(1) if shared_x_sf is not None else 0, + topk_ids.stride(0), + topk_ids.stride(1), + topk_weights.stride(0), + topk_weights.stride(1), + is_padding.stride(0) if is_padding is not None else 0, + topk_idx_out.stride(0), + topk_idx_out.stride(1), + topk_weights_out.stride(0), + topk_weights_out.stride(1), + num_tokens, + l2_activation_scale, + hidden_size, + top_k, + num_local_experts, + BLOCK_M=block_m, + BLOCK_K=128, + BLOCK_TOPK=triton.next_power_of_2(top_k), + BLOCK_E=triton.next_power_of_2(max(num_local_experts, 1)), + SHARED_BLOCK_M=shared_block_m or 1, + SHARED_MXFP8=shared_mxfp8, + num_warps=4, + ) diff --git a/vllm/model_executor/warmup/deep_gemm_warmup.py b/vllm/model_executor/warmup/deep_gemm_warmup.py index 9f7d7c7e4e..c96f1f09bf 100644 --- a/vllm/model_executor/warmup/deep_gemm_warmup.py +++ b/vllm/model_executor/warmup/deep_gemm_warmup.py @@ -14,7 +14,6 @@ from tqdm import tqdm import vllm.envs as envs from vllm.distributed.parallel_state import ( get_dp_group, - get_ep_group, is_global_first_rank, ) from vllm.forward_context import ForwardContext, override_forward_context @@ -356,7 +355,6 @@ def _get_fp8_fp4_mega_moe_m_values( dg = _import_deep_gemm() gate, experts = module.gate, module.experts uses_mega_gate = _mega_moe_uses_mega_gate(module) - num_max_tokens_per_rank = experts.get_symm_buffer().num_max_tokens_per_rank if module.use_sequence_parallel: max_tokens = cdiv(max_tokens, module.tp_size) @@ -369,14 +367,7 @@ def _get_fp8_fp4_mega_moe_m_values( if uses_mega_gate and hasattr(dg, "get_bf16_mega_gate_config") else {} ) - block_m = dg.get_block_m_for_mega_moe( - get_ep_group().world_size, - experts.num_experts, - num_max_tokens_per_rank, - num_tokens, - experts.top_k, - "fp8xfp4", - ) + block_m = experts.get_block_m(num_tokens) # Keep the largest M: DeepSeek-V4 only calls bf16_mega_gate above a # small token threshold, and a config's largest M is always above it. m_values[(tuple(gate_config.items()), block_m)] = num_tokens diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index ef24943f02..71decb5453 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -16,7 +16,6 @@ from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_world_size, ) -from vllm.distributed.eplb.eplb_state import EplbLayerState from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger from vllm.model_executor.kernels.mhc.tilelang import ( @@ -32,10 +31,9 @@ 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 ( - get_deep_gemm_mega_moe_backend, -) -from vllm.model_executor.layers.fused_moe.router.base_router import ( - eplb_map_to_physical_and_record, + MEGA_MOE_MXFP4, + DeepGemmMegaMoEExperts, + MegaMoEActivation, ) from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( fused_topk_bias, @@ -70,7 +68,6 @@ from vllm.model_executor.models.utils import ( maybe_prefix, spec_decode_needs_target_embed, ) -from vllm.model_executor.utils import set_weight_attrs from vllm.models.common.ops.sequence_parallel import ( sp_all_gather, sp_padding_mask, @@ -84,12 +81,11 @@ from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import ( DeepseekV4FlashInferSM120Attention, ) from vllm.models.deepseek_v4.nvidia.flashmla import DeepseekV4FlashMLAAttention -from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors +from vllm.utils.deep_gemm import bf16_mega_gate from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.registry import AttentionBackendEnum -from vllm.v1.worker.ubatching import dbo_current_ubatch_id from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID @@ -206,326 +202,6 @@ def make_deepseek_v4_expert_params_mapping( ] -class DeepseekV4MegaMoEExperts(nn.Module): - _symm_buffer_cache: dict[ - tuple[int, int, int, int, int, int, int, int, str], object - ] = {} - - def __init__( - self, - vllm_config: VllmConfig, - *, - num_experts: int, - num_local_experts: int, - experts_start_idx: int, - top_k: int, - 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, - ): - super().__init__() - self.prefix = prefix - self.capture_fn: Callable[[torch.Tensor], None] | None = None - self.num_experts = num_experts - self.num_local_experts = num_local_experts - self.experts_start_idx = experts_start_idx - self.experts_end_idx = experts_start_idx + num_local_experts - self.top_k = top_k - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens - - self.num_logical_experts = ( - num_logical_experts if num_logical_experts is not None else num_experts - ) - - self.eplb_state = EplbLayerState() - - 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 - ) - 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, - ) - 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_scale.quant_method = "block" - - # 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 - if prefix in compilation_config.static_forward_context: - raise ValueError(f"Duplicate layer name: {prefix}") - compilation_config.static_forward_context[prefix] = self - - def _map_global_expert_id(self, expert_id: int) -> list[int]: - """Return local (per-rank) slot offsets where logical expert - `expert_id` should land on this rank. - """ - physical_ids: list[int] = [] - for p in range(self.experts_start_idx, self.experts_end_idx): - if p % self.num_logical_experts == expert_id: - physical_ids.append(p - self.experts_start_idx) - return physical_ids - - def weight_loader( - self, - param: nn.Parameter, - loaded_weight: torch.Tensor, - weight_name: str, - shard_id: str, - expert_id: int, - return_success: bool = False, - ) -> bool | None: - local_expert_ids = self._map_global_expert_id(expert_id) - 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 - is_scale = weight_name.endswith("weight_scale") - for local_expert_id in local_expert_ids: - self.backend.load_weight( - param.data[local_expert_id], loaded_weight, shard_id, is_scale - ) - return True if return_success else None - - 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}" - - 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.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.backend - group = get_ep_group().device_group - device = torch.accelerator.current_device_index() - key = ( - id(group), - device, - self.num_experts, - self.max_num_tokens, - self.top_k, - self.hidden_size, - self.intermediate_size, - self.num_shared_experts if self.has_fused_shared_experts else 0, - backend.mma_type, - ) - symm_buffer = self._symm_buffer_cache.get(key) - if symm_buffer is None: - symm_buffer = deep_gemm.get_symm_buffer_for_mega_moe( - group, - self.num_experts, - self.max_num_tokens, - self.top_k, - self.hidden_size, - self.intermediate_size, - num_shared_experts=( - self.num_shared_experts if self.has_fused_shared_experts else 0 - ), - mma_type=backend.mma_type, - ) - self._symm_buffer_cache[key] = symm_buffer - return symm_buffer - - def set_eplb_state( - self, - moe_layer_idx: int, - expert_load_view: torch.Tensor, - logical_to_physical_map: torch.Tensor, - logical_replica_count: torch.Tensor, - ) -> None: - self.eplb_state.set_layer_state( - moe_layer_idx, - expert_load_view, - logical_to_physical_map, - logical_replica_count, - ) - - def get_expert_weights(self) -> list[torch.Tensor]: - 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 - if t.is_contiguous(): - return t.view(self.num_local_experts, -1) - elif t.dim() == 3 and t.stride(1) == 1 and t.stride(2) == t.shape[1]: - # scales have shape (E, M, N) with memory layout (E, N, M) - back = torch.transpose(t, 1, 2) - assert back.is_contiguous() - return back.view(self.num_local_experts, -1) - - raise AssertionError( - f"DSv4 EPLB {name}: non-contiguous expert tensor with " - f"unexpected layout shape={tuple(t.shape)} " - f"stride={tuple(t.stride())} dtype={t.dtype}" - ) - - return [ - _to_eplb_view("l1_packed", self._transformed_l1_weights[0]), - _to_eplb_view("l1_scale", self._transformed_l1_weights[1]), - _to_eplb_view("l2_weight", self._transformed_l2_weights[0]), - _to_eplb_view("l2_scale", self._transformed_l2_weights[1]), - ] - - def update_expert_map(self) -> None: - pass - - @property - def layer_id(self) -> int: - return extract_layer_index(self.prefix) - - def forward( - self, - hidden_states: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - *, - activation_clamp: float | None, - fast_math: bool = True, - ) -> torch.Tensor: - if hidden_states.shape[0] > self.max_num_tokens: - raise ValueError( - f"DeepSeek V4 MegaMoE got {hidden_states.shape[0]} tokens, " - f"but the symmetric buffer was sized for {self.max_num_tokens}." - ) - y = torch.empty_like(hidden_states, dtype=torch.bfloat16) - - from vllm.utils.deep_gemm import _import_deep_gemm - - deep_gemm = _import_deep_gemm() - - num_tokens = hidden_states.shape[0] - is_padding = None - if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): - is_padding = get_forward_context().is_padding - if is_padding is not None: - is_padding = is_padding[:num_tokens] - - if self.capture_fn is not None: - self.capture_fn(topk_ids) - - # EPLB: map logical expert IDs to physical replicas and record load. - eplb_state = self.eplb_state - if eplb_state.logical_to_physical_map is not None: - assert eplb_state.expert_load_view is not None - assert eplb_state.logical_replica_count is not None - assert eplb_state.should_record_tensor is not None - if is_padding is not None: - topk_ids = torch.where(is_padding.unsqueeze(1), -1, topk_ids) - topk_ids = eplb_map_to_physical_and_record( - topk_ids=topk_ids, - expert_load_view=eplb_state.expert_load_view, - logical_to_physical_map=eplb_state.logical_to_physical_map, - logical_replica_count=eplb_state.logical_replica_count, - record_enabled=eplb_state.should_record_tensor, - num_unpadded_tokens=eplb_state.num_unpadded_tokens_tensors[ - dbo_current_ubatch_id() - ] - if eplb_state.num_unpadded_tokens_tensors is not None - else None, - ) - - backend = self.backend - symm_buffer = self.get_symm_buffer() - shared_x_sf = None - shared_block_m = None - if self.has_fused_shared_experts: - shared_x_sf = symm_buffer.shared_l1_acts_sf - shared_block_m = deep_gemm.get_block_m_for_mega_moe( - get_ep_group().world_size, - self.num_experts, - symm_buffer.num_max_tokens_per_rank, - num_tokens, - self.top_k, - backend.mma_type, - ) - - prepare_megamoe_inputs( - hidden_states, - topk_weights, - topk_ids, - symm_buffer.x[:num_tokens], - symm_buffer.x_sf[:num_tokens], - symm_buffer.topk_idx[:num_tokens], - symm_buffer.topk_weights[:num_tokens], - is_padding=is_padding, - shared_x_sf=shared_x_sf, - shared_block_m=shared_block_m, - hidden_quant=backend.hidden_quant, - ) - - 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 - - -DeepseekV4MegaMoEExperts.weight_loader.supports_moe_loading = True # type: ignore[attr-defined] - - class DeepseekV4MoE(nn.Module): def __init__( self, @@ -701,7 +377,8 @@ class DeepseekV4MoE(nn.Module): and (self.use_sequence_parallel or self.tp_size == 1) ) - expert_kwargs: dict[str, typing.Any] = dict( + self.experts = DeepGemmMegaMoEExperts( + vllm_config, num_experts=self.n_physical_experts, num_local_experts=self.n_local_physical_experts, experts_start_idx=self.physical_expert_start, @@ -709,11 +386,14 @@ class DeepseekV4MoE(nn.Module): top_k=self.n_activated_experts, 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, + weight_format=MEGA_MOE_MXFP4, + activation=MegaMoEActivation( + "swiglu", + clamp=None if self.swiglu_limit is None else float(self.swiglu_limit), + ), + shared_experts=self.shared_experts if fuse_shared_experts else None, prefix=f"{prefix}.experts", ) - self.experts = DeepseekV4MegaMoEExperts(vllm_config, **expert_kwargs) if self.experts.has_fused_shared_experts: self.shared_experts = None @@ -783,13 +463,9 @@ class DeepseekV4MoE(nn.Module): return self._forward_fused_moe(hidden_states, input_ids) org_shape = hidden_states.shape - 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 ( - not mega_backend.supports_bf16_mega_gate() - or hidden_states.shape[0] <= gate_threshold - ): + if hidden_states.shape[0] <= gate_threshold: router_logits, _ = self.gate(hidden_states) topk_weights, topk_ids = fused_topk_bias( hidden_states=hidden_states, @@ -825,10 +501,10 @@ class DeepseekV4MoE(nn.Module): ] fix_routing_mask = mega_gate_metadata.hash_token_mask - topk_weights, topk_ids = mega_backend.bf16_mega_gate( - x=hidden_states, - weight=self.gate.weight, - num_topk=self.n_activated_experts, + topk_weights, topk_ids = bf16_mega_gate( + hidden_states, + self.gate.weight, + self.n_activated_experts, scoring_func=self.scoring_func, routed_scaling_factor=self.routed_scaling_factor, ep_rank=self.ep_rank, @@ -840,20 +516,9 @@ class DeepseekV4MoE(nn.Module): fix_routing_mask=fix_routing_mask, unmapped_topk_idx=unmapped_topk_idx, ) - activation_clamp = ( - float(self.swiglu_limit) if self.swiglu_limit is not None else None - ) final_hidden_states = self.experts( - hidden_states, - topk_weights, - topk_ids, - activation_clamp=activation_clamp, + hidden_states, topk_weights, topk_ids, self.shared_experts ) - - if self.shared_experts is not None: - shared_output = self.shared_experts(hidden_states) - final_hidden_states += shared_output - return final_hidden_states.view(org_shape) def _forward_fused_moe( diff --git a/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py b/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py deleted file mode 100644 index 2d93b0d430..0000000000 --- a/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py +++ /dev/null @@ -1,323 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Triton input-staging kernel for DeepSeek V4 MegaMoE. - -Quantizes hidden states to fp8 with E8M0 group scales and repacks the -routing top-k tensors into the int64/float32 layout that the DeepGEMM -MegaMoE kernels consume. -""" - -import torch - -from vllm.model_executor.layers.quantization.utils.quant_utils import ( - FP8_DTYPE, - QuantKey, - kMxfp8Dynamic, -) -from vllm.triton_utils import tl, triton - -_E8M0_SCALE_DTYPES = (torch.uint8, torch.float8_e8m0fnu) - - -@triton.jit(do_not_specialize=["num_tokens"]) -def _prepare_megamoe_inputs_kernel( - hidden_states, - x_fp8, - x_sf, - shared_x_sf, - topk_ids, - topk_weights, - is_padding, - topk_idx_out, - topk_weights_out, - hidden_stride_m: tl.constexpr, - hidden_stride_k: tl.constexpr, - x_stride_m: tl.constexpr, - x_stride_k: tl.constexpr, - x_sf_stride_m: tl.constexpr, - x_sf_stride_k: tl.constexpr, - shared_x_sf_stride_m: tl.constexpr, - shared_x_sf_stride_k: tl.constexpr, - topk_ids_stride_m: tl.constexpr, - topk_ids_stride_k: tl.constexpr, - topk_weights_stride_m: tl.constexpr, - topk_weights_stride_k: tl.constexpr, - is_padding_stride_m: tl.constexpr, - topk_idx_stride_m: tl.constexpr, - topk_idx_stride_k: tl.constexpr, - topk_weights_out_stride_m: tl.constexpr, - topk_weights_out_stride_k: tl.constexpr, - num_tokens, - hidden_size: tl.constexpr, - top_k: tl.constexpr, - BLOCK_M: tl.constexpr, - BLOCK_K: tl.constexpr, - GROUP_K: tl.constexpr, - BLOCK_TOPK: tl.constexpr, - SHARED_BLOCK_M: tl.constexpr, - USE_UE8M0: tl.constexpr, -) -> None: - token_id = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) - token_mask = token_id < num_tokens - k_block_id = tl.program_id(1) - - k_offsets = k_block_id * BLOCK_K + tl.arange(0, BLOCK_K) - k_mask = token_mask[:, None] & (k_offsets[None, :] < hidden_size) - hidden = tl.load( - hidden_states - + token_id[:, None] * hidden_stride_m - + k_offsets[None, :] * hidden_stride_k, - mask=k_mask, - other=0.0, - ).to(tl.float32) - - num_groups: tl.constexpr = BLOCK_K // GROUP_K - hidden_groups = tl.reshape(tl.abs(hidden), [BLOCK_M, num_groups, GROUP_K]) - amax = tl.max(hidden_groups, axis=2) - amax = tl.maximum(amax, 1.0e-4) - - scale = amax / 448.0 - if USE_UE8M0: - scale_bits = scale.to(tl.uint32, bitcast=True) - scale_exp = ((scale_bits >> 23) & 0xFF) + ((scale_bits & 0x7FFFFF) != 0).to( - tl.uint32 - ) - scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254) - quant_scale = (scale_exp << 23).to(tl.float32, bitcast=True) - else: - quant_scale = scale - - hidden_groups = tl.reshape(hidden, [BLOCK_M, num_groups, GROUP_K]) - scaled = hidden_groups * (1.0 / quant_scale)[:, :, None] - scaled = tl.reshape(scaled, [BLOCK_M, BLOCK_K]) - fp8 = scaled.to(tl.float8e4nv) - tl.store( - x_fp8 + token_id[:, None] * x_stride_m + k_offsets[None, :] * x_stride_k, - fp8, - mask=k_mask, - ) - - if USE_UE8M0: - scale_offsets = tl.arange(0, num_groups) - packed_scale = tl.sum(scale_exp << (scale_offsets[None, :] * 8), axis=1).to( - tl.int32 - ) - tl.store( - x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k, - packed_scale, - mask=token_mask, - ) - else: - sf_offset = tl.arange(0, num_groups) - tl.store( - x_sf - + token_id[:, None] * x_sf_stride_m - + k_block_id * x_sf_stride_k - + sf_offset[None, :], - quant_scale, - mask=token_mask[:, None], - ) - - # DeepGEMM's SM100 shared-expert TMA loads require the activation scales - # in an MN-major layout whose row permutation depends on the MegaMoE - # scheduler's runtime BLOCK_M. Write that view while the packed UE8M0 scale - # is already resident, avoiding another kernel and temporary tensor. - if USE_UE8M0 and shared_x_sf is not None: - m_block_id = token_id // SHARED_BLOCK_M - m_in_block = token_id % SHARED_BLOCK_M - aligned_block_m: tl.constexpr = triton.cdiv(SHARED_BLOCK_M, 128) * 128 - transposed_m = ( - (m_in_block // 128) * 128 + (m_in_block % 32) * 4 + (m_in_block % 128) // 32 - ) - shared_row = m_block_id * aligned_block_m + transposed_m - tl.store( - shared_x_sf - + shared_row * shared_x_sf_stride_m - + k_block_id * shared_x_sf_stride_k, - packed_scale, - mask=token_mask, - ) - - if k_block_id == 0: - topk_offsets = tl.arange(0, BLOCK_TOPK) - topk_mask = token_mask[:, None] & (topk_offsets[None, :] < top_k) - token_is_padding = tl.full((BLOCK_M,), False, tl.int1) - if is_padding is not None: - token_is_padding = tl.load( - is_padding + token_id * is_padding_stride_m, - mask=token_mask, - other=True, - ) - - ids = tl.load( - topk_ids - + token_id[:, None] * topk_ids_stride_m - + topk_offsets[None, :] * topk_ids_stride_k, - mask=topk_mask, - other=0, - ).to(tl.int64) - ids = tl.where(token_is_padding[:, None], -1, ids) - tl.store( - topk_idx_out - + token_id[:, None] * topk_idx_stride_m - + topk_offsets[None, :] * topk_idx_stride_k, - ids, - mask=topk_mask, - ) - - weights = tl.load( - topk_weights - + token_id[:, None] * topk_weights_stride_m - + topk_offsets[None, :] * topk_weights_stride_k, - mask=topk_mask, - other=0.0, - ) - weights = tl.where(token_is_padding[:, None], 0.0, weights) - tl.store( - topk_weights_out - + token_id[:, None] * topk_weights_out_stride_m - + topk_offsets[None, :] * topk_weights_out_stride_k, - weights, - mask=topk_mask, - ) - - -def _resolve_hidden_quant(hidden_quant: QuantKey, block_k: int) -> tuple[int, bool]: - """Map a hidden-state ``QuantKey`` to the kernel's ``(GROUP_K, USE_UE8M0)``.""" - scale = hidden_quant.scale - group_shape = scale.group_shape - if ( - hidden_quant.dtype != FP8_DTYPE - or not hidden_quant.symmetric - or hidden_quant.scale2 is not None - or scale.static - or not group_shape.is_per_group() - or block_k % group_shape.col != 0 - or scale.dtype not in (*_E8M0_SCALE_DTYPES, torch.float32) - ): - raise ValueError( - "DeepSeek V4 MegaMoE input staging requires symmetric dynamic fp8 " - f"per-group quantization with a group size dividing {block_k} and " - f"E8M0 or fp32 scales, got {hidden_quant}." - ) - return group_shape.col, scale.dtype in _E8M0_SCALE_DTYPES - - -def prepare_megamoe_inputs( - hidden_states: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - x_fp8: torch.Tensor, - x_sf: torch.Tensor, - topk_idx_out: torch.Tensor, - topk_weights_out: torch.Tensor, - is_padding: torch.Tensor | None = None, - shared_x_sf: torch.Tensor | None = None, - shared_block_m: int | None = None, - hidden_quant: QuantKey = kMxfp8Dynamic, -) -> None: - """Quantize hidden states and repack top-k routing for DeepGEMM MegaMoE. - - Args: - hidden_states: Input activations of shape ``[num_tokens, hidden_dim]``. - topk_weights: Router top-k weights of shape ``[num_tokens, top_k]``. - topk_ids: Router top-k expert ids of shape ``[num_tokens, top_k]``. - x_fp8: Output buffer for the fp8-quantized hidden states. - x_sf: Output buffer for the hidden-state scale factors. - topk_idx_out: Output buffer for the repacked top-k expert ids. - topk_weights_out: Output buffer for the repacked top-k weights. - is_padding: Optional per-token mask; padded tokens are not routed. - shared_x_sf: Optional output buffer for shared-expert scale factors. - Must be given together with ``shared_block_m``. - shared_block_m: Block M used to lay out ``shared_x_sf``. - hidden_quant: Hidden-state quantization scheme. E8M0 scales are packed - four per int32 into ``x_sf``; fp32 scales are stored directly. - - """ - num_tokens, hidden_size = hidden_states.shape - if num_tokens == 0: - return - if hidden_size % 128 != 0: - raise ValueError( - "DeepSeek V4 MegaMoE input staging requires hidden_size to be " - "a multiple of 128." - ) - block_k = 128 - group_k, use_ue8m0 = _resolve_hidden_quant(hidden_quant, block_k) - top_k = topk_ids.shape[1] - if topk_weights.shape != topk_ids.shape: - raise ValueError( - "DeepSeek V4 MegaMoE input staging requires topk_weights and " - "topk_ids to have the same shape." - ) - if (shared_x_sf is None) != (shared_block_m is None): - raise ValueError( - "DeepSeek V4 MegaMoE shared input staging requires both " - "shared_x_sf and shared_block_m." - ) - if shared_x_sf is not None and not use_ue8m0: - raise ValueError( - "DeepSeek V4 MegaMoE shared input staging currently requires " - "UE8M0-packed hidden scales." - ) - if shared_x_sf is not None: - assert shared_block_m is not None - if shared_block_m <= 0: - raise ValueError("MegaMoE shared_block_m must be positive.") - expected_sf_k = hidden_size // 128 - if shared_x_sf.ndim != 2 or shared_x_sf.shape[1] != expected_sf_k: - raise ValueError( - "MegaMoE shared_x_sf must have shape " - f"(*, {expected_sf_k}), got {tuple(shared_x_sf.shape)}." - ) - aligned_block_m = triton.cdiv(shared_block_m, 128) * 128 - required_rows = triton.cdiv(num_tokens, shared_block_m) * aligned_block_m - if shared_x_sf.shape[0] < required_rows: - raise ValueError( - "MegaMoE shared_x_sf has insufficient rows: requires " - f"{required_rows}, got {shared_x_sf.shape[0]}." - ) - - # On GB200, eight-row tiles win from 64 tokens; keep smaller batches untiled. - block_m = 8 if num_tokens >= 64 else 1 - grid = (triton.cdiv(num_tokens, block_m), triton.cdiv(hidden_size, block_k)) - block_topk = triton.next_power_of_2(top_k) - padding_stride_m = is_padding.stride(0) if is_padding is not None else 0 - _prepare_megamoe_inputs_kernel[grid]( - hidden_states, - x_fp8, - x_sf, - shared_x_sf, - topk_ids, - topk_weights, - is_padding, - topk_idx_out, - topk_weights_out, - hidden_states.stride(0), - hidden_states.stride(1), - x_fp8.stride(0), - x_fp8.stride(1), - x_sf.stride(0), - x_sf.stride(1), - shared_x_sf.stride(0) if shared_x_sf is not None else 0, - shared_x_sf.stride(1) if shared_x_sf is not None else 0, - topk_ids.stride(0), - topk_ids.stride(1), - topk_weights.stride(0), - topk_weights.stride(1), - padding_stride_m, - topk_idx_out.stride(0), - topk_idx_out.stride(1), - topk_weights_out.stride(0), - topk_weights_out.stride(1), - num_tokens, - hidden_size, - top_k, - BLOCK_M=block_m, - BLOCK_K=block_k, - GROUP_K=group_k, - BLOCK_TOPK=block_topk, - SHARED_BLOCK_M=shared_block_m or 1, - USE_UE8M0=use_ue8m0, - num_warps=4, - ) diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 6434ab9b48..7bf54f35d4 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -23,8 +23,10 @@ from vllm.model_executor.layers.fused_moe import ( FusedMoEFactory, fused_moe_make_expert_params_mapping, ) -from vllm.model_executor.layers.fused_moe.router.base_router import ( - eplb_map_to_physical_and_record, +from vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe import ( + MEGA_MOE_MXFP4, + DeepGemmMegaMoEExperts, + MegaMoEActivation, ) from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import ( @@ -95,8 +97,6 @@ from vllm.models.common.ops.sequence_parallel import ( sp_reduce_scatter, sp_shard, ) -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 ( LatentMoERunner, @@ -115,7 +115,6 @@ from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig from vllm.utils.math_utils import cdiv from vllm.utils.multi_stream_utils import maybe_execute_in_parallel from vllm.utils.torch_utils import aux_stream, is_meta_module -from vllm.v1.worker.ubatching import dbo_current_ubatch_id from ..common.mm_preprocess import ( KimiK3DummyInputsBuilder, @@ -330,143 +329,6 @@ class KimiRoutedOutputTransform(nn.Module): return hidden_states -class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts): - """Kimi K3 adapter for the DeepGEMM MegaMoE kernel.""" - - _kimi_symm_buffer_cache: dict[tuple[object, ...], object] = {} - _synchronized_ep_groups: set[tuple[int, int]] = set() - - def __init__( - self, - *args, - activation: str, - activation_beta: float | None, - activation_linear_beta: float | None, - **kwargs, - ): - super().__init__(*args, **kwargs) - self.activation = activation - self.activation_beta = activation_beta - self.activation_linear_beta = activation_linear_beta - - def synchronize_first_launch(self) -> None: - ep_group = get_ep_group() - device = torch.accelerator.current_device_index() - key = (id(ep_group.cpu_group), device) - if key in self._synchronized_ep_groups: - return - torch.accelerator.synchronize() - torch.distributed.barrier(group=ep_group.cpu_group) - self._synchronized_ep_groups.add(key) - - def get_symm_buffer(self): - from vllm.utils.deep_gemm import _import_deep_gemm - - deep_gemm = _import_deep_gemm() - backend = self.backend - group = get_ep_group().device_group - device = torch.accelerator.current_device_index() - key = ( - id(group), - device, - self.num_experts, - self.max_num_tokens, - self.top_k, - self.hidden_size, - self.intermediate_size, - self.activation, - backend.mma_type, - ) - symm_buffer = self._kimi_symm_buffer_cache.get(key) - if symm_buffer is None: - symm_buffer = deep_gemm.get_symm_buffer_for_mega_moe( - group, - self.num_experts, - self.max_num_tokens, - self.top_k, - self.hidden_size, - self.intermediate_size, - activation=self.activation, - mma_type=backend.mma_type, - ) - self._kimi_symm_buffer_cache[key] = symm_buffer - return symm_buffer - - def forward( - self, - hidden_states: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - *, - activation_clamp: float | None, - fast_math: bool = True, - ) -> torch.Tensor: - self.synchronize_first_launch() - if hidden_states.shape[0] > self.max_num_tokens: - raise ValueError( - f"Kimi K3 MegaMoE got {hidden_states.shape[0]} tokens, " - f"but its symmetric buffer supports {self.max_num_tokens}." - ) - y = torch.empty_like(hidden_states, dtype=torch.bfloat16) - num_tokens = hidden_states.shape[0] - is_padding = None - if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): - is_padding = get_forward_context().is_padding - if is_padding is not None: - is_padding = is_padding[:num_tokens] - - if self.capture_fn is not None: - self.capture_fn(topk_ids) - - eplb_state = self.eplb_state - if eplb_state.logical_to_physical_map is not None: - assert eplb_state.expert_load_view is not None - assert eplb_state.logical_replica_count is not None - assert eplb_state.should_record_tensor is not None - if is_padding is not None: - topk_ids = torch.where(is_padding.unsqueeze(1), -1, topk_ids) - topk_ids = eplb_map_to_physical_and_record( - topk_ids=topk_ids, - expert_load_view=eplb_state.expert_load_view, - logical_to_physical_map=eplb_state.logical_to_physical_map, - logical_replica_count=eplb_state.logical_replica_count, - record_enabled=eplb_state.should_record_tensor, - num_unpadded_tokens=eplb_state.num_unpadded_tokens_tensors[ - dbo_current_ubatch_id() - ] - if eplb_state.num_unpadded_tokens_tensors is not None - else None, - ) - - backend = self.backend - symm_buffer = self.get_symm_buffer() - prepare_megamoe_inputs( - hidden_states, - topk_weights, - topk_ids, - symm_buffer.x[:num_tokens], - symm_buffer.x_sf[:num_tokens], - symm_buffer.topk_idx[:num_tokens], - symm_buffer.topk_weights[:num_tokens], - is_padding=is_padding, - hidden_quant=backend.hidden_quant, - ) - 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, - # DeepGEMM names the SiTU gate tanh scale `activation_alpha` and the - # linear/up tanh scale `activation_beta`; beta=0 leaves up untouched. - activation=self.activation, - activation_alpha=self.activation_beta or 1.0, - activation_beta=self.activation_linear_beta or 0.0, - ) - return y - - def make_kimi_k3_mega_moe_expert_params_mapping( num_experts: int, ) -> list[tuple[str, str, int, str]]: @@ -650,7 +512,7 @@ class KimiMoE(nn.Module): f"EP size {ep_size}." ) num_local_experts = num_experts // ep_size - self.experts = KimiK3MegaMoEExperts( + self.experts = DeepGemmMegaMoEExperts( vllm_config, num_experts=num_experts, num_local_experts=num_local_experts, @@ -658,10 +520,13 @@ class KimiMoE(nn.Module): top_k=num_experts_per_token, hidden_size=self.moe_hidden_size, intermediate_size=self.padded_moe_intermediate_size, + activation=MegaMoEActivation( + "situ", + alpha=activation_situ_beta or 1.0, + beta=activation_situ_linear_beta or 0.0, + ), + weight_format=MEGA_MOE_MXFP4, prefix=f"{prefix}.experts", - activation="situ", - activation_beta=activation_situ_beta, - activation_linear_beta=activation_situ_linear_beta, ) else: self.experts = FusedMoEFactory( @@ -777,11 +642,10 @@ class KimiMoE(nn.Module): if self.use_mega_moe: assert self.routed_output_transform is not None assert topk_ids is not None + # The shared expert reads the un-projected hidden states, so it + # stays outside the latent-space MegaMoE launch. final_hidden_states = self.experts( - routed_hidden_states, - router_output, - topk_ids, - activation_clamp=None, + routed_hidden_states, router_output, topk_ids ) # The shared output is folded into the up-projection GEMM's beta-add # epilogue, so combining the two branches costs no extra kernel.