497 lines
21 KiB
Python
497 lines
21 KiB
Python
# 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)
|