feat(magi2): add fused BF16 routed MoE path (#7206)

Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>

tests/diffusion/models/magi2/test_moe_kernel_contract.py

@@ -0,0 +1,123 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+
+import vllm_omni.diffusion.models.magi2.fused_moe_kernels as kernels
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu]
+
+
+CONFIG = {
+    "BLOCK_SIZE_M": 128,
+    "BLOCK_SIZE_N": 128,
+    "BLOCK_SIZE_K": 64,
+    "GROUP_SIZE_M": 16,
+    "num_warps": 16,
+    "num_stages": 1,
+}
+
+
+class _LaunchProbe:
+    def __init__(self):
+        self.calls = []
+
+    def __getitem__(self, grid):
+        def launch(*args, **kwargs):
+            self.calls.append((grid, args, kwargs))
+
+        return launch
+
+
+def _inputs(*, gate_up: bool):
+    route_weights = torch.ones((2, 6), dtype=torch.float32)
+    if gate_up:
+        A = torch.zeros((2, 4), dtype=torch.bfloat16)
+        B = torch.zeros((2, 8, 4), dtype=torch.bfloat16)
+        C = torch.empty((12, 4), dtype=torch.bfloat16)
+        top_k = 6
+    else:
+        A = torch.zeros((12, 4), dtype=torch.bfloat16)
+        B = torch.zeros((2, 4, 4), dtype=torch.bfloat16)
+        C = torch.empty((2, 6, 4), dtype=torch.bfloat16)
+        top_k = 1
+    sorted_ids = torch.zeros(128, dtype=torch.int32)
+    expert_ids = torch.zeros(1, dtype=torch.int32)
+    padded = torch.tensor([128], dtype=torch.int32)
+    return A, B, C, route_weights, sorted_ids, expert_ids, padded, top_k
+
+
+@pytest.mark.parametrize(
+    ("gate_up", "fuse_swiglu", "mul_weight"),
+    [(True, True, False), (False, False, True)],
+)
+def test_bf16_launch_contract(monkeypatch, gate_up, fuse_swiglu, mul_weight):
+    probe = _LaunchProbe()
+    monkeypatch.setattr(kernels, "_fused_moe_bf16_kernel", probe)
+    args = _inputs(gate_up=gate_up)
+    kernels.invoke_fused_moe_bf16(*args[:-1], top_k=args[-1], config=CONFIG, fuse_swiglu=fuse_swiglu)
+    assert len(probe.calls) == 1
+    _, launch_args, launch_kwargs = probe.calls[0]
+    assert launch_args[0].dtype == torch.bfloat16
+    assert launch_kwargs["compute_type"] == kernels.tl.bfloat16
+    assert launch_kwargs["FUSE_SWIGLU"] is fuse_swiglu
+    assert launch_kwargs["MUL_ROUTED_WEIGHT"] is mul_weight
+    assert launch_kwargs["SWIGLU_ALPHA"] == (1.702 if fuse_swiglu else 0.0)
+    assert launch_kwargs["SWIGLU_LIMIT"] == (7.0 if fuse_swiglu else 0.0)
+
+
+def test_rejects_non_bf16_activation(monkeypatch):
+    args = list(_inputs(gate_up=True))
+    args[0] = args[0].float()
+    with pytest.raises(ValueError, match="requires BF16"):
+        kernels.invoke_fused_moe_bf16(*args[:-1], top_k=args[-1], config=CONFIG, fuse_swiglu=True)
+
+
+@pytest.mark.parametrize(
+    "change",
+    [
+        {"BLOCK_SIZE_N": 32},
+        {"num_warps": 3},
+        {"num_stages": 5},
+        {"unexpected": 1},
+    ],
+)
+def test_rejects_invalid_launch_config(monkeypatch, change):
+    config = dict(CONFIG)
+    config.update(change)
+    args = _inputs(gate_up=True)
+    with pytest.raises(ValueError, match="invalid MAGI-2 BF16 MoE launch config"):
+        kernels.invoke_fused_moe_bf16(*args[:-1], top_k=args[-1], config=config, fuse_swiglu=True)
+
+
+def test_rejects_incompatible_swiglu_output():
+    args = list(_inputs(gate_up=True))
+    args[2] = torch.empty((12, 8), dtype=torch.bfloat16)
+    with pytest.raises(ValueError, match="output"):
+        kernels.invoke_fused_moe_bf16(*args[:-1], top_k=args[-1], config=CONFIG, fuse_swiglu=True)
+
+
+@pytest.mark.parametrize("gate_up", [True, False])
+def test_rejects_wrong_routing_top_k(gate_up):
+    args = _inputs(gate_up=gate_up)
+    with pytest.raises(ValueError, match="routing layout"):
+        kernels.invoke_fused_moe_bf16(*args[:-1], top_k=2, config=CONFIG, fuse_swiglu=gate_up)
+
+
+@pytest.mark.parametrize(
+    ("argument", "replacement", "message"),
+    [
+        (4, torch.zeros(128, dtype=torch.float32), "sorted_token_ids"),
+        (5, torch.zeros(1, dtype=torch.float32), "expert_ids"),
+        (6, torch.tensor([128.0]), "num_tokens_post_padded"),
+        (6, torch.tensor([128, 256]), "num_tokens_post_padded"),
+        (4, torch.zeros(127, dtype=torch.int32), "block aligned"),
+        (5, torch.empty(0, dtype=torch.int32), "one expert per padded token block"),
+    ],
+)
+def test_rejects_invalid_route_metadata(argument, replacement, message):
+    args = list(_inputs(gate_up=True))
+    args[argument] = replacement
+    with pytest.raises(ValueError, match=message):
+        kernels.invoke_fused_moe_bf16(*args[:-1], top_k=args[-1], config=CONFIG, fuse_swiglu=True)