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

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

benchmarks/kernels/benchmark_magi2_bf16_moe.py

@@ -0,0 +1,310 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Compare MAGI-2 routed MoE with equal preprocessing boundaries.
+
+python -m benchmarks.kernels.benchmark_magi2_bf16_moe --tokens 2 4096 10000
+
+Default (`--mode dispatch`): the same MAGI-2 layer is timed with its native
+reference and default BF16 fast path, including routing and all preprocessing.
+With `--mode routed`, routing probabilities/IDs are inputs to BOTH timed paths. Sorting,
+padding, packing, allocations and reduction remain inside their respective
+calls. The reference uses deterministic routed storage in diagnostic mode.
+Use --mode dispatch to measure the
+actual layer with the native and default fast paths, including the router. Unsupported
+dispatch fails explicitly; it is never silently replaced by another baseline.
+
+Stage timings are standalone diagnostics, not additive or env speedup claims.
+These synthetic local shapes are not full-model performance measurements.
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import json
+import math
+import os
+import statistics
+import subprocess
+import time
+from contextlib import contextmanager
+from pathlib import Path
+from unittest.mock import patch
+
+
+def parse_args(argv=None):
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--tokens", type=int, nargs="+", default=[2, 4096])
+    parser.add_argument("--heads", type=int, default=3, help="EP4 local head count for MAGI-2 Preview.")
+    parser.add_argument("--experts", type=int, default=256, help="Released experts per head.")
+    parser.add_argument("--top-k", type=int, default=6)
+    parser.add_argument("--warmup", type=int, default=10)
+    parser.add_argument("--iters", type=int, default=50)
+    parser.add_argument("--seed", type=int, default=419)
+    parser.add_argument("--mode", choices=["routed", "dispatch"], default="dispatch")
+    parser.add_argument(
+        "--native-deterministic",
+        action="store_true",
+        help="Use the portable deterministic native reference (needed by Triton-MUSA 3.2).",
+    )
+    parser.add_argument("--stages", action="store_true")
+    parser.add_argument("--output", type=Path)
+    args = parser.parse_args(argv)
+    if min(*args.tokens, args.heads, args.experts, args.top_k, args.iters) < 1 or args.warmup < 0:
+        parser.error("Shapes/iters must be positive; warmup must be nonnegative")
+    if args.top_k > args.experts:
+        parser.error("top-k must not exceed experts per head")
+    return args
+
+
+@contextmanager
+def env_switch(enabled):
+    # Keep the legacy benchmark context API; dispatch is no longer opt-in.
+    with patch.dict(os.environ, MAGI2_DETERMINISTIC="0"):
+        yield
+
+
+def stats(values):
+    if not values or any(not math.isfinite(x) or x <= 0 for x in values):
+        raise ValueError("Missing/non-positive/non-finite timing sample")
+    ordered = sorted(values)
+    return {
+        "mean_us": statistics.mean(values),
+        "p50_us": statistics.median(values),
+        "p90_us": ordered[math.ceil(0.9 * len(ordered)) - 1],
+        "samples_us": values,
+    }
+
+
+def measure(torch, calls, args):
+    """Alternating AB/BA; cold-cache flush is outside each timed region."""
+    cuda = getattr(torch, "cuda")
+    names = list(calls)
+    sequence = names + names[::-1]
+    wall_samples = []
+    labels = []
+    counter = 0
+
+    def invoke():
+        nonlocal counter
+        name = sequence[counter % len(sequence)]
+        counter += 1
+        cuda.synchronize()
+        start = time.perf_counter_ns()
+        output = calls[name]()
+        cuda.synchronize()
+        elapsed = (time.perf_counter_ns() - start) / 1000
+        wall_samples.append(elapsed)
+        labels.append(name)
+        return output
+
+    count = len(names) * args.iters
+    cache = torch.empty(512 * 1024 * 1024, device="cuda", dtype=torch.uint8)
+    gpu_ms = []
+    for index in range(len(names) * args.warmup + count):
+        cache.zero_()
+        cuda.synchronize()
+        begin, end = cuda.Event(enable_timing=True), cuda.Event(enable_timing=True)
+        begin.record()
+        invoke()
+        end.record()
+        end.synchronize()
+        if index >= len(names) * args.warmup:
+            gpu_ms.append(begin.elapsed_time(end))
+    labels, wall_samples = labels[-count:], wall_samples[-count:]
+    result = {}
+    for name in names:
+        wall = [x for label, x in zip(labels, wall_samples) if label == name]
+        event = [float(x) * 1000 for label, x in zip(labels, gpu_ms) if label == name]
+        if len(wall) != args.iters or len(event) != args.iters:
+            raise RuntimeError("Timing harness did not return the requested sample counts")
+        result[name] = {"synchronized_wall": stats(wall), "gpu_event": stats(event)}
+    return result
+
+
+def make_inputs(torch, tokens, args):
+    # Fan-in-scaled synthetic weights. Unlike the old script, W_up and W_gate
+    # have the same distribution and top-k IDs are unique within each token.
+    torch.manual_seed(args.seed)
+    x = torch.randn(tokens, args.heads, 256, device="cuda").to(torch.bfloat16)
+    shape = (args.heads * args.experts, 256, 1280)
+    gate, up = [(torch.randn(shape, device="cuda") / 16).to(torch.bfloat16) for _ in range(2)]
+    down = (torch.randn(shape[0], 1280, 256, device="cuda") / math.sqrt(1280)).to(torch.bfloat16)
+    logits = torch.randn(args.heads, tokens, args.experts, device="cuda")
+    selected = logits.topk(args.top_k, dim=-1)
+    probabilities = selected.values.softmax(-1)
+    return x, probabilities, selected.indices, gate, up, down
+
+
+def check_and_capture(torch, moe, kernels, data, args):
+    x, probabilities, indices, gate, up, down = data
+    gather, weights, offsets = moe.global_sort_routes(probabilities, indices, args.experts)
+    sorted_routes, launches = [], []
+    scatter, invoke = kernels._deterministic_scatter, moe.invoke_fused_moe_bf16
+
+    def capture_scatter(routes, *params):
+        sorted_routes.append(routes.detach().clone())
+        return scatter(routes, *params)
+
+    def capture_launch(*params, **kwargs):
+        launches.append((params, kwargs))
+        return invoke(*params, **kwargs)
+
+    with patch.object(kernels, "_deterministic_scatter", capture_scatter):
+        moe.triton_mh_moe_forward(x, gather, weights, offsets, gate, up, down, deterministic=True)
+    packed_w13 = moe._pack_bf16_w13(gate, up)
+    route_buffers = moe._allocate_bf16_route_buffers(
+        x.shape[0] * args.heads * args.top_k,
+        args.heads * args.experts,
+        128,
+        x.device,
+    )
+    with patch.object(moe, "invoke_fused_moe_bf16", capture_launch):
+        actual = moe._bf16_fused_moe_forward(x, probabilities, indices, packed_w13, down, route_buffers)
+    if len(sorted_routes) != 1 or len(launches) != 2:
+        raise RuntimeError("Unexpected kernel dispatch: benchmark contract needs updating")
+    flat_experts = (indices + torch.arange(args.heads, device=x.device)[:, None, None] * args.experts).flatten()
+    order = flat_experts.argsort(stable=True)
+    new_routes = launches[1][0][2].reshape(-1, 256).index_select(0, order)
+    scatter_ids = gather.long() * args.heads + flat_experts[order] // args.experts
+    expected = torch.zeros(x.shape[0] * args.heads, 256, device=x.device, dtype=torch.float32)
+    expected.index_add_(0, scatter_ids, sorted_routes[0].float())
+    expected = expected.to(x.dtype).view_as(x)
+    errors = []
+    # Same expert/FP32 reduction contract as test_bf16_moe_wiring.py, rather
+    # than relaxing tolerances on the inherited BF16 scatter reduction.
+    for result, reference in ((new_routes, sorted_routes[0]), (actual, expected)):
+        if not (torch.isfinite(result).all() and torch.isfinite(reference).all()):
+            raise RuntimeError("Non-finite MoE output")
+        torch.testing.assert_close(result, reference, rtol=2e-2, atol=0.5)
+        relative_l2 = (result.float() - reference.float()).norm() / reference.float().norm().clamp_min(1e-6)
+        if relative_l2.item() >= 1e-3:
+            raise AssertionError(f"MoE relative L2 {relative_l2.item()} >= 1e-3")
+        errors.append(
+            {"max_abs": (result.float() - reference.float()).abs().max().item(), "relative_l2": relative_l2.item()}
+        )
+    return launches, errors
+
+
+def benchmark(torch, moe, kernels, tokens, args):
+    data = make_inputs(torch, tokens, args)
+    x, probabilities, indices, gate, up, down = data
+    launches, errors = check_and_capture(torch, moe, kernels, data, args)
+    packed_w13 = moe._pack_bf16_w13(gate, up)
+    route_count = tokens * args.heads * args.top_k
+    route_buffers = moe._allocate_bf16_route_buffers(route_count, args.heads * args.experts, 128, x.device)
+
+    def reference():
+        gather, weights, offsets = moe.global_sort_routes(probabilities, indices, args.experts)
+        return moe.triton_mh_moe_forward(x, gather, weights, offsets, gate, up, down, deterministic=True)
+
+    calls = {
+        "reference_deterministic_routed": reference,
+        "candidate_routed": lambda: moe._bf16_fused_moe_forward(
+            x, probabilities, indices, packed_w13, down, route_buffers
+        ),
+    }
+    if args.mode == "dispatch":
+        config = moe.Magi2MultiHeadMoEConfig(
+            args.heads * 256, args.heads, args.experts, args.top_k, 1280, torch.bfloat16
+        )
+        layer = moe.Magi2MultiHeadMoE(config).to(x.device)
+        for parameter, value in ((layer.W_gate, gate), (layer.W_up, up), (layer.W_down, down)):
+            parameter.copy_(value)
+        layer.gate.normal_(std=1 / 16)
+
+        flat_x = x.reshape(tokens, -1)
+
+        def native_dispatch():
+            x_heads = flat_x.view(tokens, args.heads, 256)
+            probs, idx = layer._route(x_heads)
+            gather, weights, offsets = moe.global_sort_routes(probs, idx, args.experts)
+            return moe.triton_mh_moe_forward(
+                x_heads,
+                gather,
+                weights,
+                offsets,
+                layer.W_gate,
+                layer.W_up,
+                layer.W_down,
+                deterministic=args.native_deterministic,
+            ).reshape(tokens, -1)
+
+        calls = {"native_dispatch": native_dispatch, "fused_dispatch": lambda: layer(flat_x)}
+        torch.testing.assert_close(calls["fused_dispatch"](), calls["native_dispatch"](), rtol=2e-2, atol=0.5)
+    timings = measure(torch, calls, args)
+    names = list(calls)
+    result = {"status": "pass", "tokens": tokens, "mode": args.mode, "correctness": errors, "timings": timings}
+    result["speedup_p50_wall"] = (
+        timings[names[0]]["synchronized_wall"]["p50_us"] / timings[names[1]]["synchronized_wall"]["p50_us"]
+    )
+    if args.stages:
+        route_ids = (indices.int() + torch.arange(args.heads, device=x.device)[:, None, None] * args.experts).reshape(
+            -1, args.top_k
+        )
+
+        def gems():
+            for params, kwargs in launches:
+                moe.invoke_fused_moe_bf16(*params, **kwargs)
+
+        result["standalone_stages_not_additive"] = measure(
+            torch,
+            {
+                "reference_sort": lambda: moe.global_sort_routes(probabilities, indices, args.experts),
+                "candidate_align": lambda: moe._align_bf16_routes(route_ids, args.heads * args.experts, 128),
+                "candidate_pack_w13": lambda: torch.stack((gate.transpose(1, 2), up.transpose(1, 2)), dim=2).reshape(
+                    args.heads * args.experts, 2560, 256
+                ),
+                "candidate_two_gemms_prepared": gems,
+            },
+            args,
+        )
+    return result
+
+
+def main():
+    args = parse_args()
+    import torch
+
+    if hasattr(torch.version, "musa") and torch.version.musa is not None:
+        importlib.import_module("torchada")  # before importing vLLM / capturing CUDA APIs
+    cuda = getattr(torch, "cuda")
+    has_gpu = cuda.is_available() or (
+        hasattr(torch.version, "musa") and torch.version.musa is not None and torch.musa.is_available()
+    )
+    if not has_gpu:
+        print(json.dumps({"status": "skip", "reason": "No visible CUDA/MUSA device"}))
+        return
+    import vllm_omni.diffusion.models.magi2.fused_moe_kernels as kernels
+    import vllm_omni.diffusion.models.magi2.mh_moe as moe
+
+    root = Path(__file__).resolve().parents[2]
+    try:
+        sha = subprocess.check_output(
+            ["git", "-C", str(root), "rev-parse", "HEAD"], text=True, stderr=subprocess.DEVNULL
+        ).strip()
+    except (OSError, subprocess.CalledProcessError):
+        sha = os.environ.get("MAGI2_BENCH_SOURCE_SHA", "unknown")
+    report = {
+        "device": cuda.get_device_name(),
+        "sm_count": cuda.get_device_properties(0).multi_processor_count,
+        "source_sha": sha,
+        "module_path": moe.__file__,
+        "torch": torch.__version__,
+        "args": {**vars(args), "output": str(args.output)},
+        "cache": "512 MiB flush outside each region",
+        "capture": "eager; JIT warmed; no graph/compile",
+        "results": [],
+    }
+    with torch.inference_mode():
+        for tokens in args.tokens:
+            result = benchmark(torch, moe, kernels, tokens, args)
+            report["results"].append(result)
+            print(json.dumps(result), flush=True)
+    if args.output:
+        args.output.parent.mkdir(parents=True, exist_ok=True)
+        args.output.write_text(json.dumps(report, indent=2) + "\n")
+
+
+if __name__ == "__main__":
+    main()

tests/benchmarks/test_magi2_bf16_moe_benchmark.py

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CPU contracts for the benchmark, independent of GPU availability."""
+
+import os
+from types import SimpleNamespace
+
+import pytest
+
+from benchmarks.kernels import benchmark_magi2_bf16_moe as bench
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_shape_sweep_and_cli_validation():
+    assert bench.parse_args(["--tokens", "2", "10000"]).tokens == [2, 10000]
+    for args in (["--iters", "0"], ["--warmup", "-1"], ["--top-k", "257"], ["--tokens", "0"]):
+        with pytest.raises(SystemExit):
+            bench.parse_args(args)
+
+
+@pytest.mark.parametrize("values", [[], [0], [-1], [float("nan")], [float("inf")]])
+def test_invalid_timings_are_not_reported(values):
+    with pytest.raises(ValueError):
+        bench.stats(values)
+
+
+def test_stats_preserve_all_samples():
+    result = bench.stats([1.0, 4.0, 2.0, 3.0])
+    assert result["p50_us"] == 2.5
+    assert result["p90_us"] == 4.0
+    assert result["samples_us"] == [1.0, 4.0, 2.0, 3.0]
+
+
+def test_switch_restores_environment_on_error(monkeypatch):
+    monkeypatch.setenv("MAGI2_DETERMINISTIC", "1")
+    with pytest.raises(RuntimeError), bench.env_switch(True):
+        assert os.environ["MAGI2_DETERMINISTIC"] == "0"
+        raise RuntimeError("compile failure")
+    assert os.environ["MAGI2_DETERMINISTIC"] == "1"
+
+
+def test_routed_reference_sort_is_inside_every_timed_call(monkeypatch):
+    data = (SimpleNamespace(device="device"), *(object() for _ in range(5)))
+    events = []
+    monkeypatch.setattr(bench, "make_inputs", lambda *args: data)
+    monkeypatch.setattr(bench, "check_and_capture", lambda *args: ([], []))
+
+    def sort(*args):
+        assert args == (data[1], data[2], 256)
+        events.append("sort")
+        return (1, 2, 3)
+
+    def legacy(*args, **kwargs):
+        assert kwargs == {"deterministic": True}
+        events.append("reference")
+
+    def candidate(*args):
+        assert args == (data[0], data[1], data[2], "packed-w13", data[5], "route-buffers")
+        events.append("candidate")
+
+    def measure(torch, calls, args):
+        assert events == ["pack", "allocate"]
+        events.clear()
+        for _ in range(2):
+            for call in calls.values():
+                call()
+        return {name: {"synchronized_wall": {"p50_us": 1}} for name in calls}
+
+    monkeypatch.setattr(bench, "measure", measure)
+
+    def pack_w13(*args):
+        events.append("pack")
+        return "packed-w13"
+
+    def allocate_route_buffers(*args):
+        events.append("allocate")
+        return "route-buffers"
+
+    moe = SimpleNamespace(
+        global_sort_routes=sort,
+        triton_mh_moe_forward=legacy,
+        _bf16_fused_moe_forward=candidate,
+        _pack_bf16_w13=pack_w13,
+        _allocate_bf16_route_buffers=allocate_route_buffers,
+    )
+    result = bench.benchmark(None, moe, None, 2, bench.parse_args(["--mode", "routed"]))
+    assert events == ["sort", "reference", "candidate"] * 2
+    assert result["mode"] == "routed"
+    assert "env_off" not in str(result["timings"])

tests/diffusion/models/magi2/test_bf16_moe_wiring.py

@@ -0,0 +1,496 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+import vllm_omni.diffusion.models.magi2.mh_moe as moe
+from vllm_omni.diffusion.models.magi2 import fused_moe_kernels
+from vllm_omni.diffusion.models.magi2.parallel import Magi2ParallelGroup
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.core_model]
+
+
+def _moe_inputs(
+    num_tokens: int,
+    num_heads: int,
+    *,
+    hidden_size: int = 7,
+    intermediate_size: int = 9,
+    num_experts: int = 4,
+    top_k: int = 2,
+    seed: int = 419,
+):
+    generator = torch.Generator().manual_seed(seed)
+    x = torch.randn(num_tokens, num_heads, hidden_size, generator=generator).to(torch.bfloat16)
+    shape = (num_heads * num_experts, hidden_size, intermediate_size)
+    gate = (torch.randn(shape, generator=generator) * 0.3).to(torch.bfloat16)
+    up = (torch.randn(shape, generator=generator) * 0.7 - 0.5).to(torch.bfloat16)
+    down = torch.randn(num_heads * num_experts, intermediate_size, hidden_size, generator=generator)
+    down = (down * 0.2).to(torch.bfloat16)
+    # Repeat head-local IDs so different heads must use different expert banks.
+    # Experts 1 and 3 have no routes, exercising empty experts and padding.
+    selected = torch.tensor([0, 2]) if top_k == 2 else torch.arange(top_k)
+    indices = selected.expand(num_heads, num_tokens, top_k).contiguous()
+    probabilities = torch.rand(num_heads, num_tokens, top_k, generator=generator)
+    probabilities = (probabilities / probabilities.sum(-1, keepdim=True)).contiguous()
+    return x, probabilities, indices, gate, up, down
+
+
+def _reference_forward(x, probabilities, indices, gate, up, down):
+    """FP32 fused-expert math with BF16 activation and weighted-route stores."""
+    num_tokens, num_heads, _ = x.shape
+    experts_per_head = gate.shape[0] // num_heads
+    output = torch.zeros_like(x)
+    for head in range(num_heads):
+        for token in range(num_tokens):
+            routes = []
+            for choice in range(indices.shape[-1]):
+                expert = head * experts_per_head + int(indices[head, token, choice])
+                gate_value = (x[token, head].float() @ gate[expert].float()).clamp(max=7)
+                up_value = (x[token, head].float() @ up[expert].float()).clamp(-7, 7)
+                hidden = (gate_value * torch.sigmoid(1.702 * gate_value) * (up_value + 1)).to(x.dtype)
+                down_value = hidden.float() @ down[expert].float()
+                routes.append((down_value * probabilities[head, token, choice]).to(x.dtype))
+            output[token, head] = torch.stack(routes).sum(0)
+    return output
+
+
+def _assert_call_args_identity(call_args, *expected):
+    assert len(call_args.args) == len(expected)
+    assert all(actual is wanted for actual, wanted in zip(call_args.args, expected))
+
+
+@pytest.mark.cpu
+def test_reference_keeps_fp32_gate_up_until_activation():
+    x = torch.ones(1, 1, 2, dtype=torch.bfloat16)
+    probabilities = torch.ones(1, 1, 1)
+    indices = torch.zeros(1, 1, 1, dtype=torch.long)
+    weight = torch.tensor([[[1.0], [2**-8]]], dtype=torch.bfloat16)
+    down = torch.ones(1, 1, 2, dtype=torch.bfloat16)
+    actual = _reference_forward(x, probabilities, indices, weight, weight, down)
+
+    projection = torch.tensor(1.0 + 2**-8)
+    expected = (projection * torch.sigmoid(1.702 * projection) * (projection + 1)).to(x.dtype)
+    rounded = projection.to(x.dtype).float()
+    early_rounded = (rounded * torch.sigmoid(1.702 * rounded) * (rounded + 1)).to(x.dtype)
+    torch.testing.assert_close(actual, expected.expand_as(actual), rtol=0, atol=0)
+    assert not torch.equal(actual, early_rounded.expand_as(actual))
+
+
+def _torch_invoke(a, b, c, weights, sorted_ids, expert_ids, num_padded, *, top_k, config, fuse_swiglu):
+    """Execute the launcher's route contract with Torch instead of Triton."""
+    route_count = weights.numel()
+    output = c.view(route_count, c.shape[-1])
+    block_size = config["BLOCK_SIZE_M"]
+    live_rows = int(num_padded.item())
+    live_blocks = live_rows // block_size
+    assert sorted_ids.numel() >= live_rows
+    assert expert_ids.numel() * block_size == sorted_ids.numel()
+    written = []
+    for block, expert in enumerate(expert_ids[:live_blocks].tolist()):
+        assert 0 <= expert < b.shape[0]
+        routes = sorted_ids[block * block_size : (block + 1) * block_size].long()
+        routes = routes[routes < route_count]
+        written.extend(routes.tolist())
+        gemm = a[routes // top_k].float() @ b[expert].float().T
+        if fuse_swiglu:
+            gate = gemm[:, 0::2].clamp(max=7)
+            up = gemm[:, 1::2].clamp(-7, 7)
+            result = gate * torch.sigmoid(1.702 * gate) * (up + 1)
+        else:
+            result = gemm * weights.reshape(-1)[routes, None]
+        output[routes] = result.to(c.dtype)
+    assert sorted(written) == list(range(route_count))
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("route_ids", [[], [2, 0, 2, 0, 2], [3] * 9 + [0, 3, 0]])
+def test_align_bf16_routes_preserves_routes_and_pads_expert_blocks(route_ids):
+    ids = torch.tensor(route_ids, dtype=torch.int64)
+    sorted_ids, expert_ids, num_padded = moe._align_bf16_routes(ids, num_experts=5, block_size=4)
+    counts = torch.bincount(ids, minlength=5)
+    expected_blocks = ((counts + 3) // 4).sum().item()
+    assert num_padded.tolist() == [expected_blocks * 4]
+    assert expert_ids.numel() >= expected_blocks
+    assert sorted_ids.dtype == expert_ids.dtype == num_padded.dtype == torch.int32
+    live_rows = int(num_padded.item())
+    assert sorted_ids.numel() >= live_rows
+    valid_ids = sorted_ids[:live_rows][sorted_ids[:live_rows] < ids.numel()].long()
+    assert sorted(valid_ids.tolist()) == list(range(ids.numel()))
+    for block, expert in enumerate(expert_ids[:expected_blocks].tolist()):
+        block_routes = sorted_ids[block * 4 : (block + 1) * 4]
+        valid = block_routes[block_routes < ids.numel()].long()
+        assert torch.all(ids[valid] == expert)
+        assert torch.all(block_routes[block_routes >= ids.numel()] == ids.numel())
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize(("num_tokens", "num_heads"), [(0, 3), (1, 1), (5, 2), (7, 3)])
+def test_bf16_forward_head_offsets_and_interleaved_weights(monkeypatch, num_tokens, num_heads):
+    inputs = _moe_inputs(num_tokens, num_heads)
+    calls = []
+
+    def invoke(*args, **kwargs):
+        calls.append(kwargs["fuse_swiglu"])
+        return _torch_invoke(*args, **kwargs)
+
+    monkeypatch.setattr(moe, "invoke_fused_moe_bf16", invoke)
+    x, probabilities, indices, gate, up, down = inputs
+    actual = moe._bf16_fused_moe_forward(x, probabilities, indices, moe._pack_bf16_w13(gate, up), down)
+    expected = _reference_forward(*inputs)
+    assert actual.shape == inputs[0].shape
+    assert actual.dtype == torch.bfloat16
+    if num_tokens:
+        assert calls == [True, False]
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    ("device_type", "dtype", "deterministic", "num_tokens", "is_musa", "expected_path"),
+    [
+        ("cuda", torch.bfloat16, None, 1, False, "triton"),
+        ("cuda", torch.bfloat16, "0", 4096, False, "bf16"),
+        ("cuda", torch.bfloat16, "1", 4096, False, "triton"),
+        ("cuda", torch.float16, None, 4096, False, "triton"),
+        ("cuda", torch.float32, None, 4096, False, "triton"),
+        ("cuda", torch.bfloat16, "0", 1, True, "bf16"),
+        ("cpu", torch.bfloat16, None, 1, False, "torch"),
+    ],
+)
+def test_local_forward_dispatch_contract(
+    monkeypatch,
+    device_type,
+    dtype,
+    deterministic,
+    num_tokens,
+    is_musa,
+    expected_path,
+):
+    config = moe.Magi2MultiHeadMoEConfig(
+        hidden_size=4,
+        num_heads=1,
+        num_experts=2,
+        top_k=1,
+        expert_intermediate_size=4,
+        params_dtype=torch.bfloat16,
+    )
+    layer = moe.Magi2MultiHeadMoE(config, ep_group=Magi2ParallelGroup(None, 1, 0))
+    if device_type == "cpu":
+        x_heads = torch.randn(num_tokens, 1, 4, dtype=dtype)
+    else:
+        # Accelerator metadata exercises dispatch without allocating on a GPU.
+        x_heads = SimpleNamespace(
+            device=SimpleNamespace(type=device_type),
+            dtype=dtype,
+            is_cuda=True,
+            shape=(num_tokens, 1, 4),
+        )
+    probabilities = torch.tensor([[[1.0]]])
+    indices = torch.tensor([[[0]]])
+    route = Mock(return_value=(probabilities, indices))
+    gather_ids, sorted_probs, offsets = object(), object(), object()
+    sort = Mock(return_value=(gather_ids, sorted_probs, offsets))
+    fast = Mock(return_value="bf16-output")
+    route_buffers = object()
+    torch_kernel = Mock(return_value="torch-output")
+    triton_kernel = Mock(return_value="triton-output")
+    monkeypatch.setattr(layer, "_route", route)
+    monkeypatch.setattr(moe, "global_sort_routes", sort)
+    monkeypatch.setattr(moe, "_bf16_fused_moe_forward", fast)
+    monkeypatch.setattr(layer, "_get_bf16_route_buffers", Mock(return_value=route_buffers))
+    monkeypatch.setattr(moe, "torch_mh_moe_forward", torch_kernel)
+    monkeypatch.setattr(moe, "triton_mh_moe_forward", triton_kernel)
+    monkeypatch.setattr(moe.current_omni_platform, "is_musa", Mock(return_value=is_musa))
+    if deterministic is None:
+        monkeypatch.delenv("MAGI2_DETERMINISTIC", raising=False)
+    else:
+        monkeypatch.setenv("MAGI2_DETERMINISTIC", deterministic)
+
+    actual = layer._local_forward(x_heads)
+
+    route.assert_called_once()
+    _assert_call_args_identity(route.call_args, x_heads)
+    if expected_path == "bf16":
+        assert actual == "bf16-output"
+        fast.assert_called_once()
+        _assert_call_args_identity(
+            fast.call_args,
+            x_heads,
+            probabilities,
+            indices,
+            layer._get_bf16_packed_w13(),
+            layer.W_down,
+            route_buffers,
+        )
+        assert fast.call_args.kwargs == {}
+        sort.assert_not_called()
+        torch_kernel.assert_not_called()
+        triton_kernel.assert_not_called()
+    else:
+        sort.assert_called_once()
+        _assert_call_args_identity(sort.call_args, probabilities, indices, layer.num_experts)
+        fast.assert_not_called()
+        if expected_path == "torch":
+            assert actual == "torch-output"
+            torch_kernel.assert_called_once()
+            _assert_call_args_identity(
+                torch_kernel.call_args,
+                x_heads,
+                gather_ids,
+                sorted_probs,
+                offsets,
+                layer.W_gate,
+                layer.W_up,
+                layer.W_down,
+            )
+            triton_kernel.assert_not_called()
+        else:
+            assert actual == "triton-output"
+            triton_kernel.assert_called_once()
+            _assert_call_args_identity(
+                triton_kernel.call_args,
+                x_heads,
+                gather_ids,
+                sorted_probs,
+                offsets,
+                layer.W_gate,
+                layer.W_up,
+                layer.W_down,
+            )
+            assert triton_kernel.call_args.kwargs == {
+                "deterministic": deterministic == "1",
+            }
+            torch_kernel.assert_not_called()
+
+
+@pytest.mark.cpu
+def test_local_forward_fast_path_propagates_kernel_failure(monkeypatch):
+    config = moe.Magi2MultiHeadMoEConfig(
+        hidden_size=4,
+        num_heads=1,
+        num_experts=2,
+        top_k=1,
+        expert_intermediate_size=4,
+        params_dtype=torch.bfloat16,
+    )
+    layer = moe.Magi2MultiHeadMoE(config, ep_group=Magi2ParallelGroup(None, 1, 0))
+    x_heads = SimpleNamespace(
+        device=SimpleNamespace(type="cuda"),
+        dtype=torch.bfloat16,
+        is_cuda=True,
+        # The fused route is selected for CUDA/MUSA only once the token count
+        # reaches the production fast-path threshold.
+        shape=(4096, 1, 4),
+    )
+    probabilities = torch.tensor([[[1.0]]])
+    indices = torch.tensor([[[0]]])
+    monkeypatch.setattr(layer, "_route", Mock(return_value=(probabilities, indices)))
+    sort = Mock()
+    torch_kernel = Mock()
+    triton_kernel = Mock()
+    error = RuntimeError("BF16 kernel failed")
+    fast = Mock(side_effect=error)
+    monkeypatch.setattr(moe, "global_sort_routes", sort)
+    monkeypatch.setattr(moe, "_bf16_fused_moe_forward", fast)
+    monkeypatch.setattr(layer, "_get_bf16_route_buffers", Mock(return_value=object()))
+    monkeypatch.setattr(moe, "torch_mh_moe_forward", torch_kernel)
+    monkeypatch.setattr(moe, "triton_mh_moe_forward", triton_kernel)
+    monkeypatch.delenv("MAGI2_DETERMINISTIC", raising=False)
+
+    with pytest.raises(RuntimeError, match="BF16 kernel failed") as exc_info:
+        layer._local_forward(x_heads)
+
+    assert exc_info.value is error
+    fast.assert_called_once()
+    sort.assert_not_called()
+    torch_kernel.assert_not_called()
+    triton_kernel.assert_not_called()
+
+
+@pytest.mark.cpu
+def test_bf16_weight_and_route_caches_are_reused_and_invalidated():
+    config = moe.Magi2MultiHeadMoEConfig(
+        hidden_size=4,
+        num_heads=1,
+        num_experts=2,
+        top_k=1,
+        expert_intermediate_size=4,
+        params_dtype=torch.bfloat16,
+    )
+    layer = moe.Magi2MultiHeadMoE(config, ep_group=Magi2ParallelGroup(None, 1, 0))
+    with torch.no_grad():
+        layer.W_gate.fill_(0.25)
+        layer.W_up.fill_(-0.5)
+
+    first = layer._get_bf16_packed_w13()
+    assert layer._get_bf16_packed_w13() is first
+    with torch.no_grad():
+        layer.W_gate.add_(1)
+    second = layer._get_bf16_packed_w13()
+    assert second is not first
+    torch.testing.assert_close(second, moe._pack_bf16_w13(layer.W_gate, layer.W_up))
+
+    buffers = layer._get_bf16_route_buffers(8, layer.W_gate.device)
+    assert layer._get_bf16_route_buffers(8, layer.W_gate.device) is buffers
+    assert layer._get_bf16_route_buffers(257, layer.W_gate.device) is not buffers
+
+
+def _gpu_device(device_type: str) -> torch.device:
+    if device_type == "musa":
+        if getattr(torch.version, "musa", None) is None or not hasattr(torch, "musa") or not torch.musa.is_available():
+            pytest.skip("MUSA runtime and device required")
+    elif not current_omni_platform.is_cuda() or current_omni_platform.get_device_count() == 0:
+        pytest.skip("CUDA runtime and device required")
+    return torch.device(device_type)
+
+
+@pytest.mark.parametrize(
+    "device_type",
+    [
+        pytest.param("cuda", marks=[pytest.mark.cuda, pytest.mark.gpu]),
+        pytest.param("musa", marks=[pytest.mark.musa, pytest.mark.gpu]),
+    ],
+)
+@pytest.mark.parametrize(
+    ("num_tokens", "num_heads", "hidden_size", "intermediate_size", "top_k", "num_experts"),
+    [
+        (1, 1, 64, 80, 2, 4),
+        (17, 3, 64, 80, 2, 4),
+        (137, 2, 64, 80, 2, 4),
+        (2, 3, 256, 1280, 6, 8),
+        (129, 1, 256, 1280, 6, 8),
+    ],
+)
+def test_bf16_forward_real_kernel_parity(
+    device_type, num_tokens, num_heads, hidden_size, intermediate_size, top_k, num_experts
+):
+    device = _gpu_device(device_type)
+    inputs = _moe_inputs(
+        num_tokens,
+        num_heads,
+        hidden_size=hidden_size,
+        intermediate_size=intermediate_size,
+        top_k=top_k,
+        num_experts=num_experts,
+    )
+    expected = _reference_forward(*inputs)
+    x, probabilities, indices, gate, up, down = (tensor.to(device) for tensor in inputs)
+    actual = moe._bf16_fused_moe_forward(x, probabilities, indices, moe._pack_bf16_w13(gate, up), down).cpu()
+    assert torch.isfinite(actual).all()
+    # BF16 activation/output rounding and cancellation across routes amplify
+    # small arithmetic differences. Bound both individual and aggregate error;
+    # retain the tighter absolute tolerance for the small masked-tile cases.
+    atol = 5e-1 if intermediate_size == 1280 else 4e-2
+    torch.testing.assert_close(actual, expected, rtol=2e-2, atol=atol)
+    relative_l2 = (actual.float() - expected.float()).norm() / expected.float().norm().clamp_min(1e-6)
+    assert relative_l2 < 1e-3
+
+
+@pytest.mark.parametrize(
+    "device_type",
+    [pytest.param("cuda", marks=pytest.mark.cuda), pytest.param("musa", marks=pytest.mark.musa)],
+)
+@pytest.mark.parametrize("seed", [419, 420])
+@pytest.mark.parametrize(
+    ("num_tokens", "num_heads", "hidden_size", "intermediate_size", "top_k", "num_experts"),
+    [
+        (1, 1, 64, 96, 2, 4),
+        (17, 3, 64, 96, 2, 4),
+        (137, 2, 64, 96, 2, 4),
+        (2, 3, 256, 1280, 6, 8),
+        (129, 1, 256, 1280, 6, 8),
+    ],
+)
+def test_bf16_forward_matches_reference_routes_and_fp32_reduction(
+    monkeypatch, device_type, seed, num_tokens, num_heads, hidden_size, intermediate_size, top_k, num_experts
+):
+    """Compare expert math and FP32 reduction, not BF16 scatter accumulation."""
+    device = _gpu_device(device_type)
+    inputs = _moe_inputs(
+        num_tokens,
+        num_heads,
+        hidden_size=hidden_size,
+        intermediate_size=intermediate_size,
+        top_k=top_k,
+        num_experts=num_experts,
+        seed=seed,
+    )
+    x, probabilities, indices, gate, up, down = (tensor.to(device) for tensor in inputs)
+    gather_ids, sorted_probs, offsets = moe.global_sort_routes(probabilities, indices, num_experts)
+    reference_routes, new_routes = [], []
+    original_scatter = fused_moe_kernels._deterministic_scatter
+    original_invoke = moe.invoke_fused_moe_bf16
+
+    def capture_reference(sorted_output, *args):
+        reference_routes.append(sorted_output.detach().cpu())
+        return original_scatter(sorted_output, *args)
+
+    def capture_new(*args, **kwargs):
+        result = original_invoke(*args, **kwargs)
+        if not kwargs.get("fuse_swiglu", False):
+            new_routes.append(args[2].detach().cpu().reshape(-1, hidden_size))
+        return result
+
+    monkeypatch.setattr(fused_moe_kernels, "_deterministic_scatter", capture_reference)
+    monkeypatch.setattr(moe, "invoke_fused_moe_bf16", capture_new)
+    # Capture before BF16 scatter; this path intentionally reduces in FP32.
+    moe.triton_mh_moe_forward(x, gather_ids, sorted_probs, offsets, gate, up, down, deterministic=True)
+    actual = moe._bf16_fused_moe_forward(x, probabilities, indices, moe._pack_bf16_w13(gate, up), down).cpu()
+    assert len(reference_routes) == len(new_routes) == 1
+    flat_experts = (inputs[2] + torch.arange(num_heads).view(-1, 1, 1) * num_experts).reshape(-1)
+    order = flat_experts.argsort(stable=True)
+    new_sorted = new_routes[0].index_select(0, order)
+    scatter_ids = gather_ids.cpu().long() * num_heads + flat_experts[order] // num_experts
+    expected_fp32 = torch.zeros(num_tokens * num_heads, hidden_size, dtype=torch.float32)
+    expected_fp32.index_add_(0, scatter_ids, reference_routes[0].float())
+    expected = expected_fp32.to(x.dtype).view(num_tokens, num_heads, hidden_size)
+    atol = 5e-1 if intermediate_size == 1280 else 4e-2
+    for result, reference in ((new_sorted, reference_routes[0]), (actual, expected)):
+        assert torch.isfinite(result).all() and torch.isfinite(reference).all()
+        torch.testing.assert_close(result, reference, rtol=2e-2, atol=atol)
+        relative_l2 = (result.float() - reference.float()).norm() / reference.float().norm().clamp_min(1e-6)
+        assert relative_l2 < 1e-3
+
+
+@pytest.mark.parametrize(
+    "device_type",
+    [
+        pytest.param("cuda", marks=[pytest.mark.cuda, pytest.mark.gpu]),
+        pytest.param("musa", marks=[pytest.mark.musa, pytest.mark.gpu]),
+    ],
+)
+def test_bf16_reference_atomic_add_capability(device_type):
+    """The non-deterministic BF16 reference path must compile on each GPU.
+
+    This is intentionally separate from candidate parity: it exercises the
+    ``deterministic=False`` branch and therefore the ``tl.atomic_add`` used by
+    the original MAGI-2 reference path.  An unsupported backend must fail the
+    test with its compiler error instead of being silently treated as a skip.
+    """
+
+    device = _gpu_device(device_type)
+    if device_type == "musa":
+        version = tuple(int(part) for part in fused_moe_kernels.triton.__version__.split(".")[:2])
+        if version < (3, 6):
+            pytest.skip("Triton-MUSA < 3.6 does not support BF16 atomic_add")
+    inputs = _moe_inputs(
+        num_tokens=2,
+        num_heads=1,
+        hidden_size=64,
+        intermediate_size=96,
+        top_k=2,
+        num_experts=4,
+        seed=419,
+    )
+    x, probabilities, indices, gate, up, down = (tensor.to(device) for tensor in inputs)
+    gather_ids, sorted_probs, offsets = moe.global_sort_routes(probabilities, indices, 4)
+    actual = moe.triton_mh_moe_forward(x, gather_ids, sorted_probs, offsets, gate, up, down, deterministic=False)
+    deterministic = moe.triton_mh_moe_forward(x, gather_ids, sorted_probs, offsets, gate, up, down, deterministic=True)
+    assert torch.isfinite(actual).all()
+    torch.testing.assert_close(actual, deterministic, rtol=2e-2, atol=4e-2)

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)

tests/kernels/test_bf16_atomic_add_minimal.py

@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""No vLLM/model/MATE dependency: python -m pytest -q this_file.py."""
+
+import importlib
+
+import pytest
+import torch
+
+if hasattr(torch.version, "musa") and torch.version.musa is not None:
+    importlib.import_module("torchada")
+
+from vllm.triton_utils import tl, triton
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.musa]
+
+
+@triton.jit
+def _atomic_sum(values, output, count: tl.constexpr, block: tl.constexpr):
+    offsets = tl.program_id(0) * block + tl.arange(0, block)
+    increments = tl.load(values + offsets, mask=offsets < count, other=0)
+    tl.atomic_add(output + offsets % 32, increments, mask=offsets < count)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_atomic_add(dtype):
+    cuda = getattr(torch, "cuda")
+    available = cuda.is_available() or (hasattr(torch, "musa") and torch.musa.is_available())
+    if not available:
+        pytest.skip("No CUDA/MUSA device")
+    if dtype == torch.bfloat16 and getattr(torch.version, "musa", None) is not None:
+        if tuple(int(part) for part in triton.__version__.split(".")[:2]) < (3, 6):
+            pytest.skip("Triton-MUSA < 3.6 does not support BF16 atomic_add")
+    # Each destination receives eight exactly representable increments from
+    # two independent blocks. Reset between runs to exercise real atomics.
+    values = torch.ones(256, device="cuda", dtype=dtype)
+    output = torch.zeros(32, device="cuda", dtype=dtype)
+    for _ in range(5):
+        output.zero_()
+        _atomic_sum[(2,)](values, output, 256, 128)
+        cuda.synchronize()
+        torch.testing.assert_close(output, torch.full_like(output, 8), rtol=0, atol=0)

vllm_omni/diffusion/models/magi2/fused_moe_kernels.py

@@ -0,0 +1,563 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# Copyright (c) 2025-2026 SandAI. All Rights Reserved.
+# ruff: noqa: N803
+
+"""BF16 routed GEMMs for the MAGI-2 fast path.
+
+The BF16 grouped GEMM is adapted from SGLang's Apache-2.0 implementation at
+``d533d3bf464d280f180e18d4d396e613a5998502``. This module implements only
+MAGI-2's interleaved gate/up GEMM with SwiGLU7 and routed down GEMM.
+The existing SandAI-derived reference kernel is retained below unchanged.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+from vllm.triton_utils import HAS_TRITON, tl, triton
+
+from vllm_omni.platforms import current_omni_platform
+
+
+def swiglu7_pair(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
+    """Reference SwiGLU7 activation used by the CPU correctness path."""
+
+    dtype = gate.dtype
+    gate = gate.float().clamp(max=7.0)
+    up = up.float().clamp(min=-7.0, max=7.0)
+    return (gate * torch.sigmoid(1.702 * gate) * (up + 1.0)).to(dtype)
+
+
+@triton.jit
+def _write_zeros_to_output(
+    c_ptr,
+    stride_cm,
+    stride_cn,
+    pid_n,
+    N,
+    offs_token,
+    token_mask,
+    BLOCK_SIZE_M,
+    BLOCK_SIZE_N,
+    compute_type,
+):
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=compute_type)
+    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+    c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]
+    tl.store(c_ptrs, accumulator, mask=token_mask[:, None] & (offs_cn[None, :] < N))
+
+
+@triton.jit
+def _fused_moe_bf16_kernel(
+    a_ptr,
+    b_ptr,
+    c_ptr,
+    topk_weights_ptr,
+    sorted_token_ids_ptr,
+    expert_ids_ptr,
+    num_tokens_post_padded_ptr,
+    N,
+    K,
+    EM,
+    num_valid_tokens,
+    stride_am,
+    stride_ak,
+    stride_be,
+    stride_bk,
+    stride_bn,
+    stride_cm,
+    stride_cn,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+    MUL_ROUTED_WEIGHT: tl.constexpr,
+    top_k: tl.constexpr,
+    compute_type: tl.constexpr,
+    even_Ks: tl.constexpr,
+    FUSE_SWIGLU: tl.constexpr,
+    SWIGLU_ALPHA: tl.constexpr,
+    SWIGLU_LIMIT: tl.constexpr,
+):
+    # Group program IDs to preserve the original L2 reuse order.
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+
+    num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
+    if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
+        return
+    offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
+    offs_token = tl.load(sorted_token_ids_ptr + offs_token_id).to(tl.int64)
+    token_mask = offs_token < num_valid_tokens
+    off_experts_i32 = tl.load(expert_ids_ptr + pid_m)
+    off_experts = off_experts_i32.to(tl.int64)
+    if off_experts == -1:
+        if FUSE_SWIGLU:
+            # The filtered down GEMM never reads these half-width rows.
+            return
+        _write_zeros_to_output(
+            c_ptr,
+            stride_cm,
+            stride_cn,
+            pid_n,
+            N,
+            offs_token,
+            token_mask,
+            BLOCK_SIZE_M,
+            BLOCK_SIZE_N,
+            compute_type,
+        )
+        return
+
+    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    a_ptrs = a_ptr + (offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak)
+    b_ptrs = b_ptr + off_experts * stride_be + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    for k_start in range(0, K, BLOCK_SIZE_K):
+        if even_Ks:
+            a = tl.load(a_ptrs, mask=token_mask[:, None], other=0.0)
+            b = tl.load(b_ptrs)
+        else:
+            a = tl.load(
+                a_ptrs,
+                mask=token_mask[:, None] & (offs_k[None, :] < K - k_start),
+                other=0.0,
+            )
+            b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k_start, other=0.0)
+        accumulator += tl.dot(a, b)
+        a_ptrs += BLOCK_SIZE_K * stride_ak
+        b_ptrs += BLOCK_SIZE_K * stride_bk
+
+    if MUL_ROUTED_WEIGHT:
+        moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0)
+        accumulator *= moe_weight[:, None]
+    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+    if FUSE_SWIGLU:
+        # Keep GEMM accumulators in FP32 through SwiGLU7.
+        acc_pairs = tl.reshape(accumulator, (BLOCK_SIZE_M, BLOCK_SIZE_N // 2, 2))
+        gate, up = tl.split(acc_pairs)
+        gate_c = tl.minimum(gate, SWIGLU_LIMIT)
+        up_c = tl.maximum(tl.minimum(up, SWIGLU_LIMIT), -SWIGLU_LIMIT)
+        out_act = (gate_c * tl.sigmoid(SWIGLU_ALPHA * gate_c) * (up_c + 1.0)).to(compute_type)
+        offs_half = pid_n * (BLOCK_SIZE_N // 2) + tl.arange(0, BLOCK_SIZE_N // 2)
+        c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_half[None, :]
+        tl.store(c_ptrs, out_act, mask=token_mask[:, None] & (offs_half[None, :] < N // 2))
+    else:
+        c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]
+        tl.store(c_ptrs, accumulator.to(compute_type), mask=token_mask[:, None] & (offs_cn[None, :] < N))
+
+
+def invoke_fused_moe_bf16(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    topk_weights: torch.Tensor,
+    sorted_token_ids: torch.Tensor,
+    expert_ids: torch.Tensor,
+    num_tokens_post_padded: torch.Tensor,
+    *,
+    top_k: int,
+    config: dict[str, int],
+    fuse_swiglu: bool,
+) -> None:
+    """Launch one BF16 MAGI-2 GEMM in the original route order.
+
+    Gate/up calls use interleaved ``B=[experts, 2*intermediate, hidden]`` and
+    retain FP32 accumulators through SwiGLU7 before writing BF16
+    ``C=[routes, intermediate]``. Down calls use ``B=[experts, hidden,
+    intermediate]`` and write ``C=[tokens, router_top_k, hidden]``. The latter
+    multiplies routing weights in FP32 before the BF16 output cast.
+    """
+    if any(t.dtype != torch.bfloat16 for t in (A, B, C)):
+        raise ValueError("MAGI-2 fused MoE requires BF16 activations, weights, and output")
+    if A.ndim != 2 or B.ndim != 3 or topk_weights.ndim != 2:
+        raise ValueError("MAGI-2 fused MoE expects A [tokens, K], B [experts, N, K], and 2D routing weights")
+    if not isinstance(fuse_swiglu, bool) or not isinstance(top_k, int) or top_k < 1:
+        raise ValueError("fuse_swiglu must be a bool and top_k must be a positive integer")
+    expected_top_k = topk_weights.shape[1] if fuse_swiglu else 1
+    if top_k != expected_top_k or A.shape[0] * top_k != topk_weights.numel():
+        raise ValueError("top_k and activation rows do not match the gate/up or down routing layout")
+    if B.shape[1] == 0 or B.shape[2] == 0 or A.shape[1] != B.shape[2]:
+        raise ValueError("MAGI-2 fused MoE requires matching nonempty GEMM dimensions")
+    if fuse_swiglu and B.shape[1] % 2:
+        raise ValueError("SwiGLU7 requires an even number of interleaved gate/up rows")
+    expected_shape = (topk_weights.numel(), B.shape[1] // 2) if fuse_swiglu else (*topk_weights.shape, B.shape[1])
+    if C.shape != expected_shape or not C.is_contiguous():
+        raise ValueError("MAGI-2 fused MoE output must have the contiguous routed output layout")
+    if not topk_weights.is_floating_point() or not topk_weights.is_contiguous():
+        raise ValueError("routing weights must be a contiguous floating-point tensor")
+    for name, tensor in (
+        ("sorted_token_ids", sorted_token_ids),
+        ("expert_ids", expert_ids),
+    ):
+        if tensor.ndim != 1 or not tensor.is_contiguous() or tensor.dtype not in (torch.int32, torch.int64):
+            raise ValueError(f"{name} must be a contiguous 1D integer tensor")
+    if num_tokens_post_padded.numel() != 1 or num_tokens_post_padded.dtype not in (
+        torch.int32,
+        torch.int64,
+    ):
+        raise ValueError("num_tokens_post_padded must contain one integer")
+    tensors = (B, C, topk_weights, sorted_token_ids, expert_ids, num_tokens_post_padded)
+    if any(t.device != A.device for t in tensors):
+        raise ValueError("MAGI-2 fused MoE tensors must share a device")
+
+    supported_config = {
+        "BLOCK_SIZE_M": (128,),
+        "BLOCK_SIZE_N": (64, 128, 256, 512),
+        "BLOCK_SIZE_K": (16, 32, 64, 128),
+        "GROUP_SIZE_M": (16,),
+        "num_warps": (4, 8, 16, 32),
+        "num_stages": (1, 2, 3, 4),
+    }
+    if config.keys() != supported_config.keys() or any(
+        type(config[key]) is not int or config[key] not in values for key, values in supported_config.items()
+    ):
+        raise ValueError("invalid MAGI-2 BF16 MoE launch config")
+    if expert_ids.numel() < triton.cdiv(sorted_token_ids.numel(), config["BLOCK_SIZE_M"]):
+        raise ValueError("expert_ids must contain one expert per padded token block")
+    if sorted_token_ids.numel() % config["BLOCK_SIZE_M"]:
+        raise ValueError("sorted_token_ids capacity must be block aligned")
+
+    def grid(meta: dict[str, int]) -> tuple[int]:
+        return (
+            triton.cdiv(sorted_token_ids.shape[0], meta["BLOCK_SIZE_M"])
+            * triton.cdiv(B.shape[1], meta["BLOCK_SIZE_N"]),
+        )
+
+    _fused_moe_bf16_kernel[grid](
+        A,
+        B,
+        C,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        B.shape[1],
+        B.shape[2],
+        sorted_token_ids.shape[0],
+        topk_weights.numel(),
+        A.stride(0),
+        A.stride(1),
+        B.stride(0),
+        B.stride(2),
+        B.stride(1),
+        C.stride(-2),
+        C.stride(-1),
+        MUL_ROUTED_WEIGHT=not fuse_swiglu,
+        top_k=top_k,
+        compute_type=tl.bfloat16,
+        even_Ks=B.shape[2] % config["BLOCK_SIZE_K"] == 0,
+        FUSE_SWIGLU=fuse_swiglu,
+        SWIGLU_ALPHA=1.702 if fuse_swiglu else 0.0,
+        SWIGLU_LIMIT=7.0 if fuse_swiglu else 0.0,
+        **config,
+    )
+
+
+def global_sort_routes(
+    topk_probs: torch.Tensor,
+    topk_indices: torch.Tensor,
+    num_experts: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Convert per-head routes into a stable flattened-expert CSR layout."""
+
+    if topk_probs.shape != topk_indices.shape or topk_indices.ndim != 3:
+        raise ValueError("top-k probabilities and indices must have the same [H,S,K] shape")
+    heads, sequence, top_k = topk_indices.shape
+    device = topk_indices.device
+    head_offset = torch.arange(heads, device=device).view(heads, 1, 1) * num_experts
+    flattened_experts = (topk_indices + head_offset).reshape(-1)
+    flat_probs = topk_probs.reshape(-1)
+    flat_tokens = torch.arange(sequence, device=device).view(1, sequence, 1).expand(heads, sequence, top_k).reshape(-1)
+    order = flattened_experts.argsort(stable=True)
+    gather_ids = flat_tokens[order].to(torch.int32)
+    sorted_probs = flat_probs[order].float()
+    counts = torch.bincount(flattened_experts, minlength=heads * num_experts)
+    offsets = torch.zeros(heads * num_experts + 1, device=device, dtype=torch.long)
+    offsets[1:] = counts.cumsum(0)
+    return gather_ids, sorted_probs, offsets
+
+
+def torch_mh_moe_forward(
+    x: torch.Tensor,
+    gather_ids: torch.Tensor,
+    probs: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    w_gate: torch.Tensor,
+    w_up: torch.Tensor,
+    w_down: torch.Tensor,
+) -> torch.Tensor:
+    """Small-shape correctness oracle for the fused expert kernel."""
+
+    if x.ndim != 3:
+        raise ValueError("multi-head MoE input must be [tokens,heads,head_dim]")
+    output = torch.zeros_like(x)
+    experts_per_head = (expert_offsets.numel() - 1) // x.shape[1]
+    for flat_expert in range(expert_offsets.numel() - 1):
+        begin = int(expert_offsets[flat_expert].item())
+        end = int(expert_offsets[flat_expert + 1].item())
+        if begin == end:
+            continue
+        head = flat_expert // experts_per_head
+        token_ids = gather_ids[begin:end].long()
+        expert_input = x.index_select(0, token_ids)[:, head]
+        gate = expert_input @ w_gate[flat_expert]
+        up = expert_input @ w_up[flat_expert]
+        hidden = swiglu7_pair(gate, up)
+        expert_output = hidden @ w_down[flat_expert]
+        expert_output = expert_output * probs[begin:end, None].to(expert_output.dtype)
+        output[:, head].index_add_(0, token_ids, expert_output)
+    return output
+
+
+# Keep these compile-time constants portable across Triton 3.2 and 3.7 while
+# allowing CPU-only imports through vLLM's no-Triton placeholder.
+_SWIGLU7_ALPHA = tl.constexpr(1.702) if HAS_TRITON else 1.702
+_SWIGLU7_LIMIT = tl.constexpr(7.0) if HAS_TRITON else 7.0
+_SWIGLU7_BIAS = tl.constexpr(1.0) if HAS_TRITON else 1.0
+
+
+@triton.jit
+def _swiglu7_kernel(gate, up, out_dtype: tl.constexpr):
+    gate_clamped = tl.minimum(gate, _SWIGLU7_LIMIT)
+    up_clamped = tl.maximum(tl.minimum(up, _SWIGLU7_LIMIT), -_SWIGLU7_LIMIT)
+    sigmoid = tl.sigmoid(_SWIGLU7_ALPHA * gate_clamped)
+    swish = gate_clamped * sigmoid
+    return (swish * (up_clamped + _SWIGLU7_BIAS)).to(out_dtype)
+
+
+@triton.jit
+def _binary_search_expert(
+    cumulative_tiles,
+    tile_id,
+    num_experts: tl.constexpr,
+    log2_num_experts: tl.constexpr,
+):
+    lo = 0
+    hi = num_experts
+    for _ in tl.static_range(0, log2_num_experts + 1):
+        mid = (lo + hi + 1) // 2
+        below = tl.load(cumulative_tiles + mid) <= tile_id
+        lo = tl.where(below, mid, lo)
+        hi = tl.where(below, hi, mid - 1)
+    return lo
+
+
+@triton.jit
+def _mh_moe_kernel(
+    x_ptr,
+    wg_ptr,
+    wu_ptr,
+    wd_ptr,
+    y_ptr,
+    gather_ids_ptr,
+    probs_ptr,
+    expert_offsets_ptr,
+    cumulative_tiles_ptr,
+    stride_x_s,
+    stride_x_h,
+    stride_x_dh,
+    stride_wg_e,
+    stride_wg_dh,
+    stride_wg_de,
+    stride_wu_e,
+    stride_wu_dh,
+    stride_wu_de,
+    stride_wd_e,
+    stride_wd_de,
+    stride_wd_dh,
+    stride_y_s,
+    stride_y_h,
+    stride_y_dh,
+    d_head: tl.constexpr,
+    d_expert: tl.constexpr,
+    num_heads: tl.constexpr,
+    num_flat_experts: tl.constexpr,
+    log2_num_experts: tl.constexpr,
+    block_t: tl.constexpr,
+    block_dh: tl.constexpr,
+    block_de: tl.constexpr,
+    acc_dtype: tl.constexpr,
+    deterministic: tl.constexpr = False,
+):
+    tile_id = tl.program_id(0)
+    total_tiles = tl.load(cumulative_tiles_ptr + num_flat_experts)
+    if tile_id >= total_tiles:
+        return
+
+    expert = _binary_search_expert(cumulative_tiles_ptr, tile_id, num_flat_experts, log2_num_experts)
+    expert_i64 = expert.to(tl.int64)
+    head = expert // (num_flat_experts // num_heads)
+    tile_in_expert = tile_id - tl.load(cumulative_tiles_ptr + expert)
+    token_start = tl.load(expert_offsets_ptr + expert) + tile_in_expert * block_t
+    expert_end = tl.load(expert_offsets_ptr + expert + 1)
+    count = tl.minimum(token_start + block_t, expert_end) - token_start
+
+    dh_block_offsets = tl.arange(0, block_dh)
+    de_block_offsets = tl.arange(0, block_de)
+    token_offsets = tl.arange(0, block_t)
+    dh_offsets = tl.arange(0, d_head)
+
+    token_positions = token_start + token_offsets
+    token_mask = token_offsets < count
+    # Token indices fit in int32, but multiplying a large packed-batch index by
+    # the hidden-width stride does not. Promote before computing element offsets.
+    gather_ids = tl.load(gather_ids_ptr + token_positions, mask=token_mask, other=0).to(tl.int64)
+    probabilities = tl.load(probs_ptr + token_positions, mask=token_mask, other=0.0)
+    x_base = gather_ids * stride_x_s + head * stride_x_h
+    output_acc = tl.zeros([block_t, d_head], dtype=acc_dtype)
+
+    for de_start in tl.range(0, d_expert, block_de):
+        de_offsets = de_start + de_block_offsets
+        gate_acc = tl.zeros([block_t, block_de], dtype=acc_dtype)
+        up_acc = tl.zeros([block_t, block_de], dtype=acc_dtype)
+        for dh_start in tl.static_range(0, d_head, block_dh):
+            local_dh = dh_start + dh_block_offsets
+            x_block = tl.load(
+                x_ptr + x_base[:, None] + local_dh[None, :] * stride_x_dh,
+                mask=token_mask[:, None],
+                other=0.0,
+            )
+            wg = tl.load(
+                wg_ptr
+                + expert_i64 * stride_wg_e
+                + local_dh[:, None] * stride_wg_dh
+                + de_offsets[None, :] * stride_wg_de
+            )
+            wu = tl.load(
+                wu_ptr
+                + expert_i64 * stride_wu_e
+                + local_dh[:, None] * stride_wu_dh
+                + de_offsets[None, :] * stride_wu_de
+            )
+            gate_acc += tl.dot(x_block, wg)
+            up_acc += tl.dot(x_block, wu)
+        hidden = _swiglu7_kernel(gate_acc, up_acc, wd_ptr.dtype.element_ty)
+        down = tl.load(
+            wd_ptr + expert_i64 * stride_wd_e + de_offsets[:, None] * stride_wd_de + dh_offsets[None, :] * stride_wd_dh
+        )
+        output_acc += tl.dot(hidden, down)
+    output_acc = output_acc * probabilities[:, None]
+
+    if deterministic:
+        output_ptrs = y_ptr + token_positions[:, None] * stride_y_s + dh_offsets[None, :] * stride_y_dh
+        tl.store(output_ptrs, output_acc.to(y_ptr.dtype.element_ty), mask=token_mask[:, None])
+    else:
+        output_base = gather_ids * stride_y_s + head * stride_y_h
+        output_ptrs = y_ptr + output_base[:, None] + dh_offsets[None, :] * stride_y_dh
+        tl.atomic_add(output_ptrs, output_acc.to(y_ptr.dtype.element_ty), mask=token_mask[:, None])
+
+
+def _deterministic_scatter(
+    sorted_output: torch.Tensor,
+    reference: torch.Tensor,
+    gather_ids: torch.Tensor,
+    expert_offsets: torch.Tensor,
+) -> torch.Tensor:
+    num_flat_experts = expert_offsets.numel() - 1
+    experts_per_head = num_flat_experts // reference.shape[1]
+    expert_lengths = torch.diff(expert_offsets)
+    head_values = torch.arange(num_flat_experts, device=gather_ids.device) // experts_per_head
+    head_ids = torch.repeat_interleave(head_values, expert_lengths)
+    scatter_ids = gather_ids.long() * reference.shape[1] + head_ids.long()
+    output = torch.zeros_like(reference).view(-1, reference.shape[-1])
+    output.scatter_add_(0, scatter_ids[:, None].expand_as(sorted_output), sorted_output.to(output.dtype))
+    return output.view_as(reference)
+
+
+def _select_block_config() -> tuple[int, int, int, int, int]:
+    """Return the reference kernel config, capped for pre-Blackwell GPUs."""
+
+    capability = current_omni_platform.get_device_capability()
+    if capability is not None and capability.major >= 10:  # Blackwell
+        return (128, 64, 32, 2, 8)
+    # BLOCK_T=128 needs 122,880 bytes of shared memory and is not safe on the
+    # qualified L20X path.  This is the reference kernel's portable config.
+    return (64, 64, 32, 2, 4)
+
+
+def triton_mh_moe_forward(
+    x: torch.Tensor,
+    gather_ids: torch.Tensor,
+    probs: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    w_gate: torch.Tensor,
+    w_up: torch.Tensor,
+    w_down: torch.Tensor,
+    *,
+    deterministic: bool = False,
+) -> torch.Tensor:
+    """Fused gather/expert/scatter kernel for released MAGI dimensions."""
+
+    routed_tokens = gather_ids.numel()
+    if routed_tokens == 0:
+        return torch.zeros_like(x)
+    d_head, d_expert = x.shape[-1], w_down.shape[1]
+    block_t, block_dh, block_de, num_stages, num_warps = _select_block_config()
+    if d_head % block_dh or d_expert % block_de:
+        return torch_mh_moe_forward(x, gather_ids, probs, expert_offsets, w_gate, w_up, w_down)
+
+    if deterministic:
+        output = torch.empty((routed_tokens, 1, d_head), device=x.device, dtype=x.dtype)
+    else:
+        output = torch.zeros_like(x)
+    num_flat_experts = expert_offsets.numel() - 1
+    expert_tiles = (torch.diff(expert_offsets) + block_t - 1) // block_t
+    cumulative_tiles = torch.cat(
+        (torch.zeros(1, dtype=torch.int32, device=x.device), expert_tiles.cumsum(0, dtype=torch.int32))
+    )
+    # Match the reference launch bound.  Empty/excess programs return after
+    # comparing against ``cumulative_tiles[-1]`` inside the kernel.
+    grid = ((routed_tokens + block_t - 1) // block_t + num_flat_experts,)
+    log2_experts = max(1, math.ceil(math.log2(max(num_flat_experts, 1) + 1)))
+    _mh_moe_kernel[grid](
+        x,
+        w_gate,
+        w_up,
+        w_down,
+        output,
+        gather_ids,
+        probs,
+        expert_offsets,
+        cumulative_tiles,
+        x.stride(0),
+        x.stride(1),
+        x.stride(2),
+        w_gate.stride(0),
+        w_gate.stride(1),
+        w_gate.stride(2),
+        w_up.stride(0),
+        w_up.stride(1),
+        w_up.stride(2),
+        w_down.stride(0),
+        w_down.stride(1),
+        w_down.stride(2),
+        output.stride(0),
+        output.stride(1),
+        output.stride(2),
+        d_head,
+        d_expert,
+        x.shape[1],
+        num_flat_experts,
+        log2_experts,
+        block_t,
+        block_dh,
+        block_de,
+        tl.float32,
+        deterministic,
+        num_stages=num_stages,
+        num_warps=num_warps,
+    )
+    if deterministic:
+        return _deterministic_scatter(output.view(routed_tokens, d_head), x, gather_ids, expert_offsets)
+    return output

vllm_omni/diffusion/models/magi2/mh_moe.py

@@ -21,22 +21,19 @@ from typing import Literal
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from vllm.triton_utils import HAS_TRITON, tl, triton
 
 from vllm_omni.platforms import current_omni_platform
 
+from .fused_moe_kernels import (
+    global_sort_routes,
+    invoke_fused_moe_bf16,
+    torch_mh_moe_forward,
+    triton_mh_moe_forward,
+)
 from .parallel import Magi2ParallelGroup, ep_dispatch, ep_undispatch, get_magi2_ep_group
 
 RoutingScore = Literal["softmax", "sigmoid"]
-
-
-def swiglu7_pair(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
-    """Released clamped SwiGLU7 expert activation, evaluated in fp32."""
-
-    dtype = gate.dtype
-    gate = gate.float().clamp(max=7.0)
-    up = up.float().clamp(min=-7.0, max=7.0)
-    return (gate * torch.sigmoid(1.702 * gate) * (up + 1.0)).to(dtype)
+_BF16_MOE_CUDA_MIN_TOKENS = 4096
 
 
 def compute_topk_probs_and_indices(
@@ -77,302 +74,127 @@ def compute_topk_probs_and_indices(
     return topk_probs, topk_indices
 
 
-def global_sort_routes(
-    topk_probs: torch.Tensor,
-    topk_indices: torch.Tensor,
+def _align_bf16_routes(
+    route_ids: torch.Tensor,
     num_experts: int,
+    block_size: int,
+    buffers: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Convert per-head routes into a stable flattened-expert CSR layout."""
+    """Build route metadata without reading the padded count back to the CPU."""
 
-    if topk_probs.shape != topk_indices.shape or topk_indices.ndim != 3:
-        raise ValueError("top-k probabilities and indices must have the same [H,S,K] shape")
-    heads, sequence, top_k = topk_indices.shape
-    device = topk_indices.device
-    head_offset = torch.arange(heads, device=device).view(heads, 1, 1) * num_experts
-    flattened_experts = (topk_indices + head_offset).reshape(-1)
-    flat_probs = topk_probs.reshape(-1)
-    flat_tokens = torch.arange(sequence, device=device).view(1, sequence, 1).expand(heads, sequence, top_k).reshape(-1)
-    order = flattened_experts.argsort(stable=True)
-    gather_ids = flat_tokens[order].to(torch.int32)
-    sorted_probs = flat_probs[order].float()
-    counts = torch.bincount(flattened_experts, minlength=heads * num_experts)
-    offsets = torch.zeros(heads * num_experts + 1, device=device, dtype=torch.long)
-    offsets[1:] = counts.cumsum(0)
-    return gather_ids, sorted_probs, offsets
+    flat_ids = route_ids.reshape(-1).to(torch.int32)
+    route_count = flat_ids.numel()
+    order = torch.argsort(flat_ids)
+    sorted_experts = flat_ids[order]
+    counts = torch.zeros(num_experts, device=flat_ids.device, dtype=torch.int64)
+    counts.scatter_add_(0, flat_ids.long(), torch.ones_like(flat_ids, dtype=torch.int64))
+    padded_counts = ((counts + block_size - 1) // block_size) * block_size
+    starts = torch.cumsum(padded_counts, 0) - padded_counts
+    ends = torch.cumsum(counts, 0)
+    begins = ends - counts
+    positions = torch.arange(route_count, device=flat_ids.device, dtype=torch.int64)
+    destinations = starts[sorted_experts.long()] + positions - begins[sorted_experts.long()]
+    if buffers is None:
+        buffers = _allocate_bf16_route_buffers(route_count, num_experts, block_size, flat_ids.device)
+    sorted_ids, expert_ids, num_padded = buffers
+    sorted_ids.fill_(route_count)
+    sorted_ids[destinations] = order.to(torch.int32)
+    block_starts = torch.arange(expert_ids.numel(), device=flat_ids.device) * block_size
+    expert_ids.copy_(torch.searchsorted(padded_counts.cumsum(0), block_starts, right=True))
+    num_padded.copy_(padded_counts.sum().reshape(1))
+    return sorted_ids, expert_ids, num_padded
 
 
-def torch_mh_moe_forward(
-    x: torch.Tensor,
-    gather_ids: torch.Tensor,
-    probs: torch.Tensor,
-    expert_offsets: torch.Tensor,
-    w_gate: torch.Tensor,
-    w_up: torch.Tensor,
-    w_down: torch.Tensor,
-) -> torch.Tensor:
-    """Small-shape correctness oracle for the fused expert kernel."""
-
-    if x.ndim != 3:
-        raise ValueError("multi-head MoE input must be [tokens,heads,head_dim]")
-    output = torch.zeros_like(x)
-    experts_per_head = (expert_offsets.numel() - 1) // x.shape[1]
-    for flat_expert in range(expert_offsets.numel() - 1):
-        begin = int(expert_offsets[flat_expert].item())
-        end = int(expert_offsets[flat_expert + 1].item())
-        if begin == end:
-            continue
-        head = flat_expert // experts_per_head
-        token_ids = gather_ids[begin:end].long()
-        expert_input = x.index_select(0, token_ids)[:, head]
-        gate = expert_input @ w_gate[flat_expert]
-        up = expert_input @ w_up[flat_expert]
-        hidden = swiglu7_pair(gate, up)
-        expert_output = hidden @ w_down[flat_expert]
-        expert_output = expert_output * probs[begin:end, None].to(expert_output.dtype)
-        output[:, head].index_add_(0, token_ids, expert_output)
-    return output
-
-
-# vLLM's no-Triton placeholder exposes ``tl.constexpr`` as ``None``.
-_SWIGLU7_ALPHA = tl.constexpr(1.702) if HAS_TRITON else 1.702
-_SWIGLU7_LIMIT = tl.constexpr(7.0) if HAS_TRITON else 7.0
-_SWIGLU7_BIAS = tl.constexpr(1.0) if HAS_TRITON else 1.0
-
-
-@triton.jit
-def _swiglu7_kernel(gate, up, out_dtype: tl.constexpr):
-    gate_clamped = tl.minimum(gate, _SWIGLU7_LIMIT)
-    up_clamped = tl.maximum(tl.minimum(up, _SWIGLU7_LIMIT), -_SWIGLU7_LIMIT)
-    sigmoid = tl.sigmoid(_SWIGLU7_ALPHA * gate_clamped)
-    swish = gate_clamped * sigmoid
-    return (swish * (up_clamped + _SWIGLU7_BIAS)).to(out_dtype)
-
-
-@triton.jit
-def _binary_search_expert(
-    cumulative_tiles,
-    tile_id,
-    num_experts: tl.constexpr,
-    log2_num_experts: tl.constexpr,
-):
-    lo = 0
-    hi = num_experts
-    for _ in tl.static_range(0, log2_num_experts + 1):
-        mid = (lo + hi + 1) // 2
-        below = tl.load(cumulative_tiles + mid) <= tile_id
-        lo = tl.where(below, mid, lo)
-        hi = tl.where(below, hi, mid - 1)
-    return lo
-
-
-@triton.jit
-def _mh_moe_kernel(
-    x_ptr,
-    wg_ptr,
-    wu_ptr,
-    wd_ptr,
-    y_ptr,
-    gather_ids_ptr,
-    probs_ptr,
-    expert_offsets_ptr,
-    cumulative_tiles_ptr,
-    stride_x_s,
-    stride_x_h,
-    stride_x_dh,
-    stride_wg_e,
-    stride_wg_dh,
-    stride_wg_de,
-    stride_wu_e,
-    stride_wu_dh,
-    stride_wu_de,
-    stride_wd_e,
-    stride_wd_de,
-    stride_wd_dh,
-    stride_y_s,
-    stride_y_h,
-    stride_y_dh,
-    d_head: tl.constexpr,
-    d_expert: tl.constexpr,
-    num_heads: tl.constexpr,
-    num_flat_experts: tl.constexpr,
-    log2_num_experts: tl.constexpr,
-    block_t: tl.constexpr,
-    block_dh: tl.constexpr,
-    block_de: tl.constexpr,
-    acc_dtype: tl.constexpr,
-    deterministic: tl.constexpr = False,
-):
-    tile_id = tl.program_id(0)
-    total_tiles = tl.load(cumulative_tiles_ptr + num_flat_experts)
-    if tile_id >= total_tiles:
-        return
-
-    expert = _binary_search_expert(cumulative_tiles_ptr, tile_id, num_flat_experts, log2_num_experts)
-    expert_i64 = expert.to(tl.int64)
-    head = expert // (num_flat_experts // num_heads)
-    tile_in_expert = tile_id - tl.load(cumulative_tiles_ptr + expert)
-    token_start = tl.load(expert_offsets_ptr + expert) + tile_in_expert * block_t
-    expert_end = tl.load(expert_offsets_ptr + expert + 1)
-    count = tl.minimum(token_start + block_t, expert_end) - token_start
-
-    dh_block_offsets = tl.arange(0, block_dh)
-    de_block_offsets = tl.arange(0, block_de)
-    token_offsets = tl.arange(0, block_t)
-    dh_offsets = tl.arange(0, d_head)
-
-    token_positions = token_start + token_offsets
-    token_mask = token_offsets < count
-    # Token indices fit in int32, but multiplying a large packed-batch index by
-    # the hidden-width stride does not. Promote before computing element offsets.
-    gather_ids = tl.load(gather_ids_ptr + token_positions, mask=token_mask, other=0).to(tl.int64)
-    probabilities = tl.load(probs_ptr + token_positions, mask=token_mask, other=0.0)
-    x_base = gather_ids * stride_x_s + head * stride_x_h
-    output_acc = tl.zeros([block_t, d_head], dtype=acc_dtype)
-
-    for de_start in tl.range(0, d_expert, block_de):
-        de_offsets = de_start + de_block_offsets
-        gate_acc = tl.zeros([block_t, block_de], dtype=acc_dtype)
-        up_acc = tl.zeros([block_t, block_de], dtype=acc_dtype)
-        for dh_start in tl.static_range(0, d_head, block_dh):
-            local_dh = dh_start + dh_block_offsets
-            x_block = tl.load(
-                x_ptr + x_base[:, None] + local_dh[None, :] * stride_x_dh,
-                mask=token_mask[:, None],
-                other=0.0,
-            )
-            wg = tl.load(
-                wg_ptr
-                + expert_i64 * stride_wg_e
-                + local_dh[:, None] * stride_wg_dh
-                + de_offsets[None, :] * stride_wg_de
-            )
-            wu = tl.load(
-                wu_ptr
-                + expert_i64 * stride_wu_e
-                + local_dh[:, None] * stride_wu_dh
-                + de_offsets[None, :] * stride_wu_de
-            )
-            gate_acc += tl.dot(x_block, wg)
-            up_acc += tl.dot(x_block, wu)
-        hidden = _swiglu7_kernel(gate_acc, up_acc, wd_ptr.dtype.element_ty)
-        down = tl.load(
-            wd_ptr + expert_i64 * stride_wd_e + de_offsets[:, None] * stride_wd_de + dh_offsets[None, :] * stride_wd_dh
-        )
-        output_acc += tl.dot(hidden, down)
-    output_acc = output_acc * probabilities[:, None]
-
-    if deterministic:
-        output_ptrs = y_ptr + token_positions[:, None] * stride_y_s + dh_offsets[None, :] * stride_y_dh
-        tl.store(output_ptrs, output_acc.to(y_ptr.dtype.element_ty), mask=token_mask[:, None])
-    else:
-        output_base = gather_ids * stride_y_s + head * stride_y_h
-        output_ptrs = y_ptr + output_base[:, None] + dh_offsets[None, :] * stride_y_dh
-        tl.atomic_add(output_ptrs, output_acc.to(y_ptr.dtype.element_ty), mask=token_mask[:, None])
-
-
-def _deterministic_scatter(
-    sorted_output: torch.Tensor,
-    reference: torch.Tensor,
-    gather_ids: torch.Tensor,
-    expert_offsets: torch.Tensor,
-) -> torch.Tensor:
-    num_flat_experts = expert_offsets.numel() - 1
-    experts_per_head = num_flat_experts // reference.shape[1]
-    expert_lengths = torch.diff(expert_offsets)
-    head_values = torch.arange(num_flat_experts, device=gather_ids.device) // experts_per_head
-    head_ids = torch.repeat_interleave(head_values, expert_lengths)
-    scatter_ids = gather_ids.long() * reference.shape[1] + head_ids.long()
-    output = torch.zeros_like(reference).view(-1, reference.shape[-1])
-    output.scatter_add_(0, scatter_ids[:, None].expand_as(sorted_output), sorted_output.to(output.dtype))
-    return output.view_as(reference)
-
-
-def _select_block_config() -> tuple[int, int, int, int, int]:
-    """Return the reference kernel config, capped for pre-Blackwell GPUs."""
-
-    capability = current_omni_platform.get_device_capability()
-    if capability is not None and capability.major >= 10:  # Blackwell
-        return (128, 64, 32, 2, 8)
-    # BLOCK_T=128 needs 122,880 bytes of shared memory and is not safe on the
-    # qualified L20X path.  This is the reference kernel's portable config.
-    return (64, 64, 32, 2, 4)
-
-
-def triton_mh_moe_forward(
-    x: torch.Tensor,
-    gather_ids: torch.Tensor,
-    probs: torch.Tensor,
-    expert_offsets: torch.Tensor,
-    w_gate: torch.Tensor,
-    w_up: torch.Tensor,
-    w_down: torch.Tensor,
-    *,
-    deterministic: bool = False,
-) -> torch.Tensor:
-    """Fused gather/expert/scatter kernel for released MAGI dimensions."""
-
-    routed_tokens = gather_ids.numel()
-    if routed_tokens == 0:
-        return torch.zeros_like(x)
-    d_head, d_expert = x.shape[-1], w_down.shape[1]
-    block_t, block_dh, block_de, num_stages, num_warps = _select_block_config()
-    if d_head % block_dh or d_expert % block_de:
-        return torch_mh_moe_forward(x, gather_ids, probs, expert_offsets, w_gate, w_up, w_down)
-
-    if deterministic:
-        output = torch.empty((routed_tokens, 1, d_head), device=x.device, dtype=x.dtype)
-    else:
-        output = torch.zeros_like(x)
-    num_flat_experts = expert_offsets.numel() - 1
-    expert_tiles = (torch.diff(expert_offsets) + block_t - 1) // block_t
-    cumulative_tiles = torch.cat(
-        (torch.zeros(1, dtype=torch.int32, device=x.device), expert_tiles.cumsum(0, dtype=torch.int32))
+def _allocate_bf16_route_buffers(
+    route_count: int, num_experts: int, block_size: int, device: torch.device
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    # Each expert adds at most block_size - 1 padding rows. The kernel uses
+    # the device-side padded count to skip programs beyond the live prefix.
+    capacity = math.ceil((route_count + num_experts * (block_size - 1)) / block_size) * block_size
+    return (
+        torch.empty(capacity, device=device, dtype=torch.int32),
+        torch.empty(capacity // block_size, device=device, dtype=torch.int32),
+        torch.empty(1, device=device, dtype=torch.int32),
     )
-    # Match the reference launch bound.  Empty/excess programs return after
-    # comparing against ``cumulative_tiles[-1]`` inside the kernel.
-    grid = ((routed_tokens + block_t - 1) // block_t + num_flat_experts,)
-    log2_experts = max(1, math.ceil(math.log2(max(num_flat_experts, 1) + 1)))
-    _mh_moe_kernel[grid](
-        x,
-        w_gate,
-        w_up,
-        w_down,
-        output,
-        gather_ids,
-        probs,
-        expert_offsets,
-        cumulative_tiles,
-        x.stride(0),
-        x.stride(1),
-        x.stride(2),
-        w_gate.stride(0),
-        w_gate.stride(1),
-        w_gate.stride(2),
-        w_up.stride(0),
-        w_up.stride(1),
-        w_up.stride(2),
-        w_down.stride(0),
-        w_down.stride(1),
-        w_down.stride(2),
-        output.stride(0),
-        output.stride(1),
-        output.stride(2),
-        d_head,
-        d_expert,
-        x.shape[1],
-        num_flat_experts,
-        log2_experts,
-        block_t,
-        block_dh,
-        block_de,
-        tl.float32,
-        deterministic,
-        num_stages=num_stages,
-        num_warps=num_warps,
+
+
+def _pack_bf16_w13(w_gate: torch.Tensor, w_up: torch.Tensor) -> torch.Tensor:
+    """Pack gate/up weights into adjacent rows for the grouped W13 GEMM."""
+
+    num_experts, hidden_size, intermediate_size = w_gate.shape
+    return torch.stack((w_gate.transpose(1, 2), w_up.transpose(1, 2)), dim=2).reshape(
+        num_experts, 2 * intermediate_size, hidden_size
     )
-    if deterministic:
-        return _deterministic_scatter(output.view(routed_tokens, d_head), x, gather_ids, expert_offsets)
-    return output
+
+
+def _bf16_fused_moe_forward(
+    x_heads: torch.Tensor,
+    probabilities: torch.Tensor,
+    indices: torch.Tensor,
+    packed_w13: torch.Tensor,
+    w_down: torch.Tensor,
+    route_buffers: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
+) -> torch.Tensor:
+    """Run the BF16 grouped GEMM path with head-local routing."""
+
+    num_tokens, num_heads, hidden_size = x_heads.shape
+    top_k = probabilities.shape[-1]
+    num_experts = packed_w13.shape[0]
+    if num_experts % num_heads:
+        raise ValueError("expert bank must be divisible by local MoE heads")
+    experts_per_head = num_experts // num_heads
+    if num_tokens == 0:
+        return torch.zeros_like(x_heads)
+    head_offsets = (
+        torch.arange(num_heads, device=x_heads.device, dtype=torch.int32).view(num_heads, 1, 1) * experts_per_head
+    )
+    route_ids = (indices.to(torch.int32) + head_offsets).reshape(num_heads * num_tokens, top_k)
+    route_weights = probabilities.reshape(num_heads * num_tokens, top_k).contiguous()
+    sorted_ids, expert_ids, num_padded = _align_bf16_routes(route_ids, num_experts, 128, route_buffers)
+
+    hidden = x_heads.permute(1, 0, 2).contiguous().reshape(num_heads * num_tokens, hidden_size)
+    intermediate_size = packed_w13.shape[1] // 2
+    intermediate = torch.empty(
+        (num_heads * num_tokens * top_k, intermediate_size), device=x_heads.device, dtype=x_heads.dtype
+    )
+    # Keep the qualified MUSA point unchanged.  CUDA/H20 benefits from the
+    # pre-Blackwell tile found by the MAGI-2 BF16 sweep (smaller K/warp count
+    # and deeper pipelining reduce register pressure).  The launch contract
+    # remains the same; only legal, device-specific values are selected.
+    config = {
+        "BLOCK_SIZE_M": 128,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 32,
+        "GROUP_SIZE_M": 16,
+        "num_warps": 4 if not current_omni_platform.is_musa() else 16,
+        "num_stages": 3 if not current_omni_platform.is_musa() else 1,
+    }
+    invoke_fused_moe_bf16(
+        hidden,
+        packed_w13,
+        intermediate,
+        route_weights,
+        sorted_ids,
+        expert_ids,
+        num_padded,
+        top_k=top_k,
+        config=config,
+        fuse_swiglu=True,
+    )
+    route_output = torch.empty((num_heads * num_tokens, top_k, hidden_size), device=x_heads.device, dtype=x_heads.dtype)
+    invoke_fused_moe_bf16(
+        intermediate,
+        w_down.transpose(1, 2),
+        route_output,
+        route_weights,
+        sorted_ids,
+        expert_ids,
+        num_padded,
+        top_k=1,
+        config=config,
+        fuse_swiglu=False,
+    )
+    return route_output.sum(dim=1).reshape(num_heads, num_tokens, hidden_size).permute(1, 0, 2)
 
 
 @dataclass(frozen=True)
@@ -450,6 +272,59 @@ class Magi2MultiHeadMoE(nn.Module):
             parameter = getattr(target, parts[-1])
             parameter.mmap_weight_transform = self.ep_slice
 
+        # Derived, non-persistent weight storage.  The model loader populates
+        # it after checkpoint loading; the lazy guard also supports direct
+        # layer construction and later weight reloads in tests/tools.
+        self.register_buffer("_bf16_packed_w13", None, persistent=False)
+        self._bf16_packed_w13_key: tuple | None = None
+        self._bf16_route_buffers: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None
+
+    def _apply(self, fn, recurse: bool = True):
+        # Do not migrate an extra full-size derived weight bank on .to().
+        self._bf16_packed_w13 = None
+        self._bf16_packed_w13_key = None
+        self._bf16_route_buffers = None
+        return super()._apply(fn, recurse=recurse)
+
+    def _get_bf16_packed_w13(self) -> torch.Tensor:
+        # In-place loading increments _version; mmap/.data replacement may
+        # not. Include object/storage identity and handle inference tensors,
+        # which intentionally do not expose a version counter.
+        key = tuple(
+            (id(weight), weight.data_ptr(), None if weight.is_inference() else weight._version)
+            for weight in (self.W_gate, self.W_up)
+        )
+        packed = self._bf16_packed_w13
+        if (
+            packed is None
+            or self._bf16_packed_w13_key != key
+            or packed.device != self.W_gate.device
+            or packed.dtype != self.W_gate.dtype
+        ):
+            # [E, I, 2, D] -> [E, 2I, D]: adjacent rows are gate/up pairs.
+            with torch.no_grad():
+                packed = _pack_bf16_w13(self.W_gate, self.W_up)
+            self._bf16_packed_w13 = packed
+            self._bf16_packed_w13_key = key
+        return packed
+
+    def prepare_bf16_weights(self) -> None:
+        """Materialize derived BF16 weights outside the inference hot path."""
+
+        # Reload may mutate inference tensors without a version counter.
+        self._bf16_packed_w13 = None
+        self._bf16_packed_w13_key = None
+        if self.W_gate.dtype == torch.bfloat16 and self.W_gate.device.type not in ("cpu", "meta"):
+            self._get_bf16_packed_w13()
+
+    def _get_bf16_route_buffers(self, route_count: int, device: torch.device):
+        buffers = self._bf16_route_buffers
+        capacity = math.ceil((route_count + self.local_flatten_num_experts * 127) / 128) * 128
+        if buffers is None or buffers[0].numel() != capacity or buffers[0].device != device:
+            buffers = _allocate_bf16_route_buffers(route_count, self.local_flatten_num_experts, 128, device)
+            self._bf16_route_buffers = buffers
+        return buffers
+
     def ep_slice(self, checkpoint_tensor: torch.Tensor) -> torch.Tensor:
         """Slice flattened ``(head,expert)`` checkpoint rows for this rank."""
 
@@ -492,6 +367,21 @@ class Magi2MultiHeadMoE(nn.Module):
 
     def _local_forward(self, x_heads: torch.Tensor) -> torch.Tensor:
         probabilities, indices = self._route(x_heads)
+        if (
+            x_heads.device.type in ("cuda", "musa")
+            and x_heads.dtype == torch.bfloat16
+            and os.environ.get("MAGI2_DETERMINISTIC", "0") != "1"
+            and (current_omni_platform.is_musa() or x_heads.shape[0] >= _BF16_MOE_CUDA_MIN_TOKENS)
+            and not torch.compiler.is_compiling()
+        ):
+            return _bf16_fused_moe_forward(
+                x_heads,
+                probabilities,
+                indices,
+                self._get_bf16_packed_w13(),
+                self.W_down,
+                self._get_bf16_route_buffers(x_heads.shape[0] * self.local_num_heads * self.top_k, x_heads.device),
+            )
         gather_ids, sorted_probs, offsets = global_sort_routes(probabilities, indices, self.num_experts)
         if x_heads.is_cuda:
             return triton_mh_moe_forward(

vllm_omni/diffusion/models/magi2/modeling_magi2.py

@@ -806,6 +806,10 @@ class Magi2PreviewTransformer(nn.Module):
                 module.router.expert_bias.copy_(module.router.expert_bias_ema)
                 loaded.add(f"{module_name}.router.expert_bias")
 
+        for module in self.modules():
+            if isinstance(module, Magi2MultiHeadMoE):
+                module.prepare_bf16_weights()
+
         missing = set(targets) - loaded
         if missing:
             raise ValueError(f"MAGI-2 preview checkpoint is missing {len(missing)} weights: {sorted(missing)[:8]}")