[MoE] Unify DeepGEMM MegaMoE experts behind one feature-matrix layer; add NVFP4 (#61538)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Co-authored-by: Kevin H. Luu <khluu000@gmail.com>
This commit is contained in:
Yongye Zhu
co-authored by
Claude Opus 5.5
Kevin H. Luu
Claude Opus 5.5
Kevin H. Luu
parent
4d9a8d6a41
commit
4d58fa254a
@@ -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:
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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]
|
||||
|
||||
@@ -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),
|
||||
)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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,
|
||||
)
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user