From 29b818a78503bd43a79923bd83ad1fa948d3a9b2 Mon Sep 17 00:00:00 2001 From: R0CKSTAR Date: Mon, 28 Sep 2026 23:20:30 +0800 Subject: [PATCH] feat(magi2): add fused BF16 routed MoE path (#7206) Signed-off-by: Xiaodong Ye --- .../kernels/benchmark_magi2_bf16_moe.py | 310 ++++++++++ .../test_magi2_bf16_moe_benchmark.py | 90 +++ .../models/magi2/test_bf16_moe_wiring.py | 496 +++++++++++++++ .../models/magi2/test_moe_kernel_contract.py | 123 ++++ tests/kernels/test_bf16_atomic_add_minimal.py | 43 ++ .../models/magi2/fused_moe_kernels.py | 563 ++++++++++++++++++ vllm_omni/diffusion/models/magi2/mh_moe.py | 488 ++++++--------- .../diffusion/models/magi2/modeling_magi2.py | 4 + 8 files changed, 1818 insertions(+), 299 deletions(-) create mode 100644 benchmarks/kernels/benchmark_magi2_bf16_moe.py create mode 100644 tests/benchmarks/test_magi2_bf16_moe_benchmark.py create mode 100644 tests/diffusion/models/magi2/test_bf16_moe_wiring.py create mode 100644 tests/diffusion/models/magi2/test_moe_kernel_contract.py create mode 100644 tests/kernels/test_bf16_atomic_add_minimal.py create mode 100644 vllm_omni/diffusion/models/magi2/fused_moe_kernels.py diff --git a/benchmarks/kernels/benchmark_magi2_bf16_moe.py b/benchmarks/kernels/benchmark_magi2_bf16_moe.py new file mode 100644 index 000000000..4a1ca7ccf --- /dev/null +++ b/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() diff --git a/tests/benchmarks/test_magi2_bf16_moe_benchmark.py b/tests/benchmarks/test_magi2_bf16_moe_benchmark.py new file mode 100644 index 000000000..bbc00594d --- /dev/null +++ b/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"]) diff --git a/tests/diffusion/models/magi2/test_bf16_moe_wiring.py b/tests/diffusion/models/magi2/test_bf16_moe_wiring.py new file mode 100644 index 000000000..f0a1a98f7 --- /dev/null +++ b/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) diff --git a/tests/diffusion/models/magi2/test_moe_kernel_contract.py b/tests/diffusion/models/magi2/test_moe_kernel_contract.py new file mode 100644 index 000000000..da017c2de --- /dev/null +++ b/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) diff --git a/tests/kernels/test_bf16_atomic_add_minimal.py b/tests/kernels/test_bf16_atomic_add_minimal.py new file mode 100644 index 000000000..f252a0f80 --- /dev/null +++ b/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) diff --git a/vllm_omni/diffusion/models/magi2/fused_moe_kernels.py b/vllm_omni/diffusion/models/magi2/fused_moe_kernels.py new file mode 100644 index 000000000..4031daebb --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/models/magi2/mh_moe.py b/vllm_omni/diffusion/models/magi2/mh_moe.py index 9e1e0788a..f504ed458 100644 --- a/vllm_omni/diffusion/models/magi2/mh_moe.py +++ b/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( diff --git a/vllm_omni/diffusion/models/magi2/modeling_magi2.py b/vllm_omni/diffusion/models/magi2/modeling_magi2.py index 0009a2ecc..1223043a7 100644 --- a/vllm_omni/diffusion/models/magi2/modeling_magi2.py +++ b/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]}")