anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
This commit is contained in:
R0CKSTAR authored and GitHub committed 2026-09-28 23:20:30 +08:00
1 parent 79a46a4f54
commit 29b818a785
8 files changed
+1818 -299
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/models/magi2/test_bf16_moe_wiring.py
+496
View File
@@ -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
+123
View File
@@ -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)