R0CKSTARandHongsheng Liu
3bf96080f3
Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
375 lines
16 KiB
Python
375 lines
16 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.layers import mhc
|
|
from vllm_omni.diffusion.models.magi2 import layers as model_layers
|
|
|
|
pytestmark = [pytest.mark.diffusion, pytest.mark.core_model]
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def clear_mhc_failure_cache():
|
|
"""Keep the process-lifetime production cache isolated between tests."""
|
|
|
|
mhc._FAILED_MHC_KERNELS.clear()
|
|
mhc._WARNED_MHC_KERNELS.clear()
|
|
yield
|
|
mhc._FAILED_MHC_KERNELS.clear()
|
|
mhc._WARNED_MHC_KERNELS.clear()
|
|
|
|
|
|
def coefficients(tokens=7, streams=4):
|
|
torch.manual_seed(71)
|
|
# Actual model split views: post and residual share a wider FP32 row.
|
|
packed = torch.randn(tokens, streams * (streams + 2))
|
|
post = packed[:, streams : 2 * streams]
|
|
residual = packed[:, 2 * streams :].view(tokens, streams, streams)
|
|
return post, residual, torch.tensor([0.8]), torch.randn(streams), torch.tensor([1.1]), torch.randn(streams, streams)
|
|
|
|
|
|
def reference(args, *, scale=0.125, iterations=20, epsilon=1e-12, out_dtype=torch.float32):
|
|
post, residual, ap, bp, ar, br = args
|
|
post = 2.0 * torch.sigmoid(ap * scale * post + bp.unsqueeze(0))
|
|
residual = ar * scale * residual.float() + br.unsqueeze(0).float()
|
|
residual = torch.exp(residual - residual.amax(dim=(-2, -1), keepdim=True))
|
|
for _ in range(iterations):
|
|
residual = residual / (residual.sum(-2, keepdim=True) + epsilon)
|
|
residual = residual / (residual.sum(-1, keepdim=True) + epsilon)
|
|
return post.to(out_dtype), residual.to(out_dtype)
|
|
|
|
|
|
def mix_reference(args):
|
|
streams, output, post, matrix = args
|
|
branch = torch.einsum("tn,tc->tnc", post, output)
|
|
mixed = torch.einsum("tij,tjc->tic", matrix, streams)
|
|
return mixed + branch
|
|
|
|
|
|
def mix_inputs(tokens=7, hidden=70, dtype=torch.float32, streams=4):
|
|
torch.manual_seed(41)
|
|
return (
|
|
torch.randn(tokens, streams, hidden).to(dtype),
|
|
torch.randn(tokens, hidden).to(dtype),
|
|
torch.rand(tokens, streams).to(dtype),
|
|
torch.randn(tokens, streams, streams).to(dtype),
|
|
)
|
|
|
|
|
|
@pytest.mark.cpu
|
|
@pytest.mark.parametrize("streams", [3, 4])
|
|
@pytest.mark.parametrize("iterations", [-1, 0, 1, 20, 65])
|
|
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
|
|
def test_native_coefficient_contract(streams, iterations, dtype):
|
|
args = coefficients(streams=streams)
|
|
kwargs = dict(scale=0.125, iterations=iterations, out_dtype=dtype)
|
|
actual = mhc.MHCPostResidual().forward_cuda(*args, **kwargs)
|
|
for a, b in zip(actual, reference(args, **kwargs)):
|
|
torch.testing.assert_close(a, b, rtol=0, atol=0)
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_native_gradients_and_reused_parameters():
|
|
op = mhc.MHCPostResidual()
|
|
args = tuple(x.detach().clone().requires_grad_() for x in coefficients())
|
|
expected_args = tuple(x.detach().clone().requires_grad_() for x in args)
|
|
result = op.forward_cuda(*args, scale=0.125, out_dtype=torch.float32)
|
|
expected = reference(expected_args)
|
|
sum(x.square().sum() for x in result).backward()
|
|
sum(x.square().sum() for x in expected).backward()
|
|
for a, b in zip(args, expected_args):
|
|
torch.testing.assert_close(a.grad, b.grad, rtol=0, atol=0)
|
|
with torch.no_grad():
|
|
args[3].add_(2)
|
|
result = op.forward_cuda(*args, scale=0.125, out_dtype=torch.float32)
|
|
for a, b in zip(result, reference(args)):
|
|
torch.testing.assert_close(a, b, rtol=0, atol=0)
|
|
|
|
|
|
@pytest.mark.cpu
|
|
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
|
|
def test_native_mix_dtype_and_noncontiguous(dtype):
|
|
args = mix_inputs(dtype=dtype)
|
|
args = (args[0][..., ::2], args[1][..., ::2], *args[2:])
|
|
torch.testing.assert_close(mhc.MHCMix().forward_cuda(*args), mix_reference(args), rtol=0, atol=0)
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_meta_shapes_and_empty():
|
|
args = tuple(x.to("meta") for x in coefficients())
|
|
post, matrix = mhc.MHCPostResidual().forward_cuda(*args, scale=0.125, out_dtype=torch.bfloat16)
|
|
assert post.shape == (7, 4) and matrix.shape == (7, 4, 4)
|
|
assert post.dtype == matrix.dtype == torch.bfloat16
|
|
mix = mhc.MHCMix().forward_cuda(*[x.to("meta") for x in mix_inputs()])
|
|
assert mix.shape == (7, 4, 70) and mix.device.type == "meta"
|
|
args = coefficients(tokens=0)
|
|
post, matrix = mhc.MHCPostResidual().forward_cuda(*args, scale=0.125, out_dtype=torch.float32)
|
|
assert post.shape == (0, 4) and matrix.shape == (0, 4, 4)
|
|
assert mhc.MHCMix().forward_cuda(*mix_inputs(tokens=0)).shape == (0, 4, 70)
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_sinkhorn_alias_and_rectangular_native_contract():
|
|
assert model_layers.sinkhorn_knopp is mhc.sinkhorn_knopp
|
|
logits = torch.randn(2, 3, 4)
|
|
expected = torch.exp(logits - logits.amax((-2, -1), keepdim=True))
|
|
torch.testing.assert_close(mhc.sinkhorn_knopp(logits, 0, 1e-12), expected, rtol=0, atol=0)
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_model_wiring_uses_fused_eager_path(monkeypatch):
|
|
calls = []
|
|
|
|
class Prepare:
|
|
def __call__(self, *args, **kwargs):
|
|
calls.append("dispatch")
|
|
return reference(args, **kwargs)
|
|
|
|
def forward_native(self, *args, **kwargs):
|
|
calls.append("native")
|
|
return reference(args, **kwargs)
|
|
|
|
monkeypatch.setattr(model_layers, "_mhc_post_residual", Prepare())
|
|
handler = model_layers.MHCHandler(4, 16)
|
|
post, residual, ap, bp, ar, br = coefficients()
|
|
handler.compute_post_residual((ap, bp, post), (ar, br, residual), out_dtype=torch.float32)
|
|
assert calls == ["dispatch"]
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_compile_keeps_native_graph():
|
|
torch._dynamo.reset()
|
|
targets: list[str] = []
|
|
|
|
def backend(gm, _inputs):
|
|
targets.extend(str(node.target) for node in gm.graph.nodes)
|
|
return gm.forward
|
|
|
|
op = torch.compile(mhc.MHCPostResidual(), fullgraph=True, backend=backend)
|
|
args = coefficients()
|
|
actual = op(*args, scale=0.125, iterations=2, out_dtype=torch.float32)
|
|
for a, b in zip(actual, reference(args, iterations=2)):
|
|
torch.testing.assert_close(a, b, rtol=0, atol=0)
|
|
assert any("sigmoid" in x for x in targets) and any("sum" in x for x in targets)
|
|
assert not any("triton" in x for x in targets)
|
|
|
|
|
|
DEVICES = [pytest.param("cuda", marks=pytest.mark.cuda), pytest.param("musa", marks=pytest.mark.musa)]
|
|
|
|
|
|
def device(name):
|
|
api = getattr(torch, name, None)
|
|
if api is None or not api.is_available():
|
|
pytest.skip(f"requires {name}")
|
|
return torch.device(name)
|
|
|
|
|
|
def spy(monkeypatch, name):
|
|
original = getattr(mhc, name)
|
|
launches = []
|
|
|
|
class Spy:
|
|
def __getitem__(self, grid):
|
|
launches.append(grid)
|
|
return original[grid]
|
|
|
|
monkeypatch.setattr(mhc, name, Spy())
|
|
return launches
|
|
|
|
|
|
@pytest.mark.parametrize("device_name", DEVICES)
|
|
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
|
|
@pytest.mark.parametrize("tokens,iterations", [(0, 20), (1, 0), (17, 20)])
|
|
def test_coefficients_gpu(monkeypatch, device_name, dtype, tokens, iterations):
|
|
dev = device(device_name)
|
|
# Preserve the model's strided split views on the GPU too.
|
|
args = coefficients(tokens)
|
|
packed = torch.zeros(tokens, 24, device=dev)
|
|
packed[:, 4:8] = args[0].to(dev)
|
|
packed[:, 8:] = args[1].reshape(tokens, 16).to(dev)
|
|
args = (packed[:, 4:8], packed[:, 8:].view(tokens, 4, 4), *(x.to(dev) for x in args[2:]))
|
|
launches = spy(monkeypatch, "_mhc_post_residual_kernel")
|
|
actual = mhc.MHCPostResidual()(*args, scale=0.125, iterations=iterations, out_dtype=dtype)
|
|
getattr(torch, device_name).synchronize()
|
|
assert len(launches) == bool(tokens)
|
|
tolerance = {torch.float32: 1e-5, torch.float16: 2e-3, torch.bfloat16: 1.6e-2}[dtype]
|
|
for a, b in zip(actual, reference(args, iterations=iterations, out_dtype=dtype)):
|
|
torch.testing.assert_close(a, b, rtol=tolerance, atol=1e-6)
|
|
|
|
|
|
@pytest.mark.parametrize("device_name", DEVICES)
|
|
def test_sinkhorn_gpu_shift_and_extremes(device_name):
|
|
dev = device(device_name)
|
|
args = tuple(x.to(dev) for x in coefficients())
|
|
op = mhc.MHCPostResidual()
|
|
post, matrix = op(*args, scale=1.0, out_dtype=torch.float32)
|
|
shifted = (args[0], args[1] + 16.0, *args[2:])
|
|
shifted_post, shifted_matrix = op(*shifted, scale=1.0, out_dtype=torch.float32)
|
|
torch.testing.assert_close(shifted_matrix, matrix, rtol=1e-5, atol=3e-6)
|
|
torch.testing.assert_close(shifted_post, post, rtol=0, atol=0)
|
|
torch.testing.assert_close(matrix.sum(-1), torch.ones_like(post), rtol=1e-5, atol=1e-5)
|
|
# A well-conditioned case converges in both dimensions at the requested
|
|
# iteration count. Extreme matrices below need not converge in 20 steps.
|
|
balanced_args = list(args)
|
|
balanced_args[1] = args[1] * 0.125
|
|
balanced_args[5] = torch.zeros_like(args[5])
|
|
_, balanced = op(*balanced_args, scale=1.0, out_dtype=torch.float32)
|
|
for dim in (-2, -1):
|
|
torch.testing.assert_close(balanced.sum(dim), torch.ones_like(post), rtol=1e-5, atol=1e-5)
|
|
extreme = list(args)
|
|
extreme[1] = torch.eye(4, device=dev).expand(7, -1, -1).contiguous() * 1000
|
|
extreme[1][0, 0, 0] = float("nan")
|
|
extreme[1][1, 0, 0] = float("inf")
|
|
for iters in (0, 20):
|
|
actual = op(*extreme, scale=1.0, iterations=iters, out_dtype=torch.float32)
|
|
expected = reference(extreme, scale=1.0, iterations=iters)
|
|
for a, b in zip(actual, expected):
|
|
torch.testing.assert_close(a, b, rtol=1e-5, atol=1e-5, equal_nan=True)
|
|
|
|
|
|
@pytest.mark.parametrize("device_name", DEVICES)
|
|
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
|
|
@pytest.mark.parametrize("tokens,hidden", [(0, 70), (2, 70), (17, 3072)])
|
|
def test_mix_gpu(monkeypatch, device_name, dtype, tokens, hidden):
|
|
dev = device(device_name)
|
|
args = tuple(x.to(dev) for x in mix_inputs(tokens, hidden, dtype))
|
|
launches = spy(monkeypatch, "_mhc_mix_kernel")
|
|
actual = mhc.MHCMix()(*args)
|
|
getattr(torch, device_name).synchronize()
|
|
assert len(launches) == bool(tokens)
|
|
tolerance = {torch.float32: 1e-5, torch.float16: 2e-3, torch.bfloat16: 1.6e-2}[dtype]
|
|
torch.testing.assert_close(actual, mix_reference(args), rtol=tolerance, atol=1e-5)
|
|
|
|
|
|
@pytest.mark.parametrize("device_name", DEVICES)
|
|
def test_mix_preserves_separate_bf16_rounding(device_name):
|
|
dev = device(device_name)
|
|
streams = torch.ones(1, 4, 1, device=dev, dtype=torch.bfloat16)
|
|
branch = torch.full((1, 1), 1.0078125, device=dev, dtype=torch.bfloat16)
|
|
post = torch.full((1, 4), 1.03125, device=dev, dtype=torch.bfloat16)
|
|
matrix = torch.eye(4, device=dev, dtype=torch.bfloat16).unsqueeze(0)
|
|
args = (streams, branch, post, matrix)
|
|
expected = mix_reference(args)
|
|
once = (
|
|
torch.einsum("tij,tjc->tic", matrix.float(), streams.float())
|
|
+ torch.einsum("tn,tc->tnc", post.float(), branch.float())
|
|
).bfloat16()
|
|
assert not torch.equal(expected, once), "fixture must distinguish the cast boundaries"
|
|
torch.testing.assert_close(mhc.MHCMix()(*args), expected, rtol=0, atol=0)
|
|
|
|
|
|
@pytest.mark.parametrize("device_name", DEVICES)
|
|
def test_gpu_unsupported_layout_falls_back(monkeypatch, device_name):
|
|
dev = device(device_name)
|
|
|
|
class Forbidden:
|
|
def __getitem__(self, _grid):
|
|
pytest.fail("unsupported layouts must not launch the kernel")
|
|
|
|
monkeypatch.setattr(mhc, "_mhc_post_residual_kernel", Forbidden())
|
|
monkeypatch.setattr(mhc, "_mhc_mix_kernel", Forbidden())
|
|
args = tuple(x.to(dev) for x in coefficients(streams=3))
|
|
actual = mhc.MHCPostResidual()(*args, scale=0.125, out_dtype=torch.float32)
|
|
for a, b in zip(actual, reference(args)):
|
|
torch.testing.assert_close(a, b, rtol=0, atol=0)
|
|
args = tuple(x.to(dev) for x in mix_inputs())
|
|
args = (args[0][..., ::2], args[1][..., ::2], *args[2:])
|
|
torch.testing.assert_close(mhc.MHCMix()(*args), mix_reference(args), rtol=0, atol=0)
|
|
# A singleton alpha with extra rank changes broadcasting, so it cannot be
|
|
# treated like the scalar alpha loaded by the eager kernel.
|
|
broadcast_args = list(x.to(dev) for x in coefficients())
|
|
broadcast_args[2] = broadcast_args[2].view(1, 1, 1)
|
|
actual = mhc.MHCPostResidual()(*broadcast_args, scale=0.125, out_dtype=torch.float32)
|
|
for a, b in zip(actual, reference(broadcast_args)):
|
|
torch.testing.assert_close(a, b, rtol=0, atol=0)
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_no_triton_uses_native_fallback(monkeypatch):
|
|
monkeypatch.setattr(mhc, "HAS_TRITON", False)
|
|
args = coefficients(tokens=2)
|
|
actual = mhc.MHCPostResidual().forward_cuda(*args, scale=0.125, out_dtype=torch.float32)
|
|
for got, expected in zip(actual, reference(args)):
|
|
torch.testing.assert_close(got, expected, rtol=0, atol=0)
|
|
|
|
mix_args = mix_inputs(tokens=2)
|
|
torch.testing.assert_close(mhc.MHCMix().forward_cuda(*mix_args), mix_reference(mix_args), rtol=0, atol=0)
|
|
|
|
|
|
@pytest.mark.cpu
|
|
@pytest.mark.parametrize(
|
|
("device_type", "native_required"),
|
|
[("cpu", True), ("meta", True), ("hip", True), ("xpu", True), ("cuda", False), ("musa", False)],
|
|
)
|
|
def test_unvalidated_backend_uses_native(device_type, native_required):
|
|
tensor = SimpleNamespace(device=SimpleNamespace(type=device_type), requires_grad=False)
|
|
assert mhc._native_required((tensor,)) is native_required
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_kernel_failure_is_cached_and_falls_back(monkeypatch):
|
|
monkeypatch.setattr(mhc, "HAS_TRITON", True)
|
|
monkeypatch.setattr(mhc, "_native_required", lambda tensors: False)
|
|
warnings = []
|
|
monkeypatch.setattr(mhc.logger, "warning", lambda *args, **kwargs: warnings.append((args, kwargs)))
|
|
|
|
class Broken:
|
|
def __getitem__(self, _grid):
|
|
def launch(*args, **kwargs):
|
|
raise RuntimeError("synthetic kernel failure")
|
|
|
|
return launch
|
|
|
|
monkeypatch.setattr(mhc, "_mhc_post_residual_kernel", Broken())
|
|
args = tuple(x.contiguous() for x in coefficients(tokens=2))
|
|
op = mhc.MHCPostResidual()
|
|
expected = reference(args)
|
|
for _ in range(2):
|
|
actual = op.forward_cuda(*args, scale=0.125, out_dtype=torch.float32)
|
|
for got, want in zip(actual, expected):
|
|
torch.testing.assert_close(got, want, rtol=0, atol=0)
|
|
assert ("cpu", 0, "post_residual") in mhc._FAILED_MHC_KERNELS
|
|
assert len(warnings) == 1
|
|
assert "post_residual" in str(warnings[0])
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_kernel_oom_is_propagated_without_caching(monkeypatch):
|
|
monkeypatch.setattr(mhc, "HAS_TRITON", True)
|
|
monkeypatch.setattr(mhc, "_native_required", lambda tensors: False)
|
|
|
|
class OutOfMemory:
|
|
def __getitem__(self, _grid):
|
|
def launch(*args, **kwargs):
|
|
raise RuntimeError("synthetic CUDA out of memory")
|
|
|
|
return launch
|
|
|
|
monkeypatch.setattr(mhc, "_mhc_post_residual_kernel", OutOfMemory())
|
|
args = tuple(x.contiguous() for x in coefficients(tokens=2))
|
|
with pytest.raises(RuntimeError, match="out of memory"):
|
|
mhc.MHCPostResidual().forward_cuda(*args, scale=0.125, out_dtype=torch.float32)
|
|
assert ("cpu", 0, "post_residual") not in mhc._FAILED_MHC_KERNELS
|
|
|
|
|
|
@pytest.mark.parametrize("device_name", DEVICES)
|
|
def test_gpu_compiled_parameters_do_not_get_cached(monkeypatch, device_name):
|
|
dev = device(device_name)
|
|
|
|
class Forbidden:
|
|
def __getitem__(self, _grid):
|
|
pytest.fail("compilation must retain the native expression")
|
|
|
|
monkeypatch.setattr(mhc, "_mhc_post_residual_kernel", Forbidden())
|
|
op = torch.compile(mhc.MHCPostResidual(), backend="eager", fullgraph=True)
|
|
args = list(x.to(dev) for x in coefficients())
|
|
for bias in (0.0, 2.0, -1.0):
|
|
args[3].fill_(bias)
|
|
actual = op(*args, scale=0.125, out_dtype=torch.float32)
|
|
for a, b in zip(actual, reference(args)):
|
|
torch.testing.assert_close(a, b, rtol=0, atol=0)
|