From 3bf96080f305fa91e775c01c52c07dabc7ef8a25 Mon Sep 17 00:00:00 2001 From: R0CKSTAR Date: Fri, 25 Sep 2026 16:25:12 +0800 Subject: [PATCH] feat(magi2): add fused mHC post-processing (#7261) Signed-off-by: Xiaodong Ye Co-authored-by: Hongsheng Liu --- benchmarks/kernels/_magi2_gpu_bench.py | 194 +++++++++++ benchmarks/kernels/benchmark_magi2_mhc.py | 141 ++++++++ tests/diffusion/layers/test_mhc.py | 374 +++++++++++++++++++++ vllm_omni/diffusion/layers/mhc.py | 279 +++++++++++++++ vllm_omni/diffusion/models/magi2/layers.py | 36 +- 5 files changed, 1007 insertions(+), 17 deletions(-) create mode 100644 benchmarks/kernels/_magi2_gpu_bench.py create mode 100644 benchmarks/kernels/benchmark_magi2_mhc.py create mode 100644 tests/diffusion/layers/test_mhc.py create mode 100644 vllm_omni/diffusion/layers/mhc.py diff --git a/benchmarks/kernels/_magi2_gpu_bench.py b/benchmarks/kernels/_magi2_gpu_bench.py new file mode 100644 index 000000000..cdfde092e --- /dev/null +++ b/benchmarks/kernels/_magi2_gpu_bench.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Cold-cache torch.cuda timing and correctness gates for MAGI-2 microbenchmarks.""" + +import argparse +import importlib +import importlib.metadata +import json +import logging +import math +import os +import subprocess +import sys +from collections.abc import Callable +from pathlib import Path +from types import ModuleType +from typing import Any + +LOGGER = logging.getLogger(__name__) + + +def emit(**fields: Any) -> None: + LOGGER.info(json.dumps(fields, sort_keys=True)) + + +def add_arguments(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--tokens", type=int, nargs="+", default=[1, 4, 16, 64, 129, 4096]) + parser.add_argument("--dtype", choices=["bfloat16", "float16", "float32"], default="bfloat16") + parser.add_argument("--num-tests", type=int, default=30, help="Measured samples per implementation.") + parser.add_argument("--warmup", type=int, default=10, help="Interleaved warmup calls before sampling.") + parser.add_argument("--seed", type=int, default=42) + + +def load_musa(args: argparse.Namespace, parser: argparse.ArgumentParser) -> Any: + if min(*args.tokens, args.num_tests) <= 0 or args.warmup < 0: + parser.error("tokens and num-tests must be positive; warmup must be nonnegative") + logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stdout) + torch = importlib.import_module("torch") + cuda = getattr(torch, "cuda") + # torchada is needed only on MUSA; after it is loaded, use torch.cuda for + # both CUDA and MUSA as documented by torchada. + if hasattr(torch.version, "musa") and torch.version.musa is not None: + try: + importlib.import_module("torchada") + except ModuleNotFoundError as error: + if error.name != "torchada": + raise + 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: + emit(status="skip", reason="CUDA/MUSA device is unavailable", torch=torch.__version__) + return None + torch.manual_seed(args.seed) + cuda.manual_seed_all(args.seed) + return torch + + +def report_environment(torch: Any, args: argparse.Namespace, module: ModuleType) -> None: + root = Path(__file__).resolve().parents[2] + packages = {} + for name in ("torch", "torch_musa", "torchada", "triton", "vllm", "vllm-omni"): + try: + packages[name] = importlib.metadata.version(name) + except importlib.metadata.PackageNotFoundError: + packages[name] = None + cuda = getattr(torch, "cuda") + properties = cuda.get_device_properties(cuda.current_device()) + try: + source_sha = subprocess.check_output(["git", "-C", str(root), "rev-parse", "HEAD"], text=True).strip() + source_dirty = bool(subprocess.check_output(["git", "-C", str(root), "status", "--porcelain"], text=True)) + except (subprocess.CalledProcessError, FileNotFoundError): + source_sha = os.environ.get("MAGI2_BENCH_SOURCE_SHA", "unknown") + source_dirty = None + emit( + status="environment", + command=sys.argv, + source_root=str(root), + source_sha=source_sha, + source_dirty=source_dirty, + kernel_source=module.__file__, + device=cuda.get_device_name(cuda.current_device()), + sm_count=getattr(properties, "multi_processor_count", None), + packages=packages, + env={name: os.environ.get(name) for name in ("MUSA_VISIBLE_DEVICES", "TORCHDYNAMO_DISABLE")}, + config=vars(args), + mode="eager operator diagnostic; no compile or graph capture", + comparison="forward_native versus the default CustomOp dispatch", + timing="torch.cuda.Event; alternating native/fused calls; vanilla 512 MiB flush outside timed regions", + timing_scope="whole callable, including device launch gaps; not isolated kernel duration", + l2_flush=True, + l2_flush_bytes=512 * 1024 * 1024, + ) + + +class _LaunchProbe: + def __init__(self, kernel: Any) -> None: + self.kernel = kernel + self.count = 0 + + def __getitem__(self, grid: Any) -> Any: + self.count += 1 + return self.kernel[grid] + + +def check_pair( + torch: Any, + reference: Callable, + fused: Callable, + module: ModuleType, + kernels: tuple[str, ...], + *, + rtol: float, + atol: float, +) -> dict: + cuda = getattr(torch, "cuda") + expected = reference() + probes = {name: _LaunchProbe(getattr(module, name)) for name in kernels} + try: + for name, probe in probes.items(): + setattr(module, name, probe) + actual = fused() + cuda.synchronize() + finally: + for name, probe in probes.items(): + setattr(module, name, probe.kernel) + if any(probe.count != 1 for probe in probes.values()): + raise RuntimeError(f"Expected one launch per fused kernel; got {[(k, p.count) for k, p in probes.items()]}") + expected = expected if isinstance(expected, tuple) else (expected,) + actual = actual if isinstance(actual, tuple) else (actual,) + errors = [] + for result, baseline in zip(actual, expected, strict=True): + for value in (baseline, result): + if not torch.isfinite(value).all().item(): + raise AssertionError("Nonfinite benchmark output") + if value.numel() > 1 and (value == value.flatten()[0]).all().item(): + raise AssertionError("Constant benchmark output") + torch.testing.assert_close(result, baseline, rtol=rtol, atol=atol) + difference = (result.float() - baseline.float()).abs() + errors.append( + { + "max_abs": difference.max().item(), + "max_rel": (difference / baseline.float().abs().clamp_min(1e-6)).max().item(), + } + ) + return {"parity": "pass", "errors": errors, "rtol": rtol, "atol": atol, "kernels": list(kernels)} + + +def _percentile(samples: list[float], quantile: float) -> float: + ordered = sorted(samples) + position = (len(ordered) - 1) * quantile + lower, upper = math.floor(position), math.ceil(position) + return ordered[lower] + (ordered[upper] - ordered[lower]) * (position - lower) + + +def measure_pair(torch: Any, reference: Callable, fused: Callable, args: argparse.Namespace) -> dict: + cuda = getattr(torch, "cuda") + functions = (reference, fused) + order: list[int] = [] + + def alternate() -> Any: + mode = len(order) % 2 + order.append(mode) + return functions[mode]() + + times_ms = [] + for _ in range(args.warmup): + alternate() + cuda.synchronize() + order.clear() + flush = torch.empty(512 * 1024 * 1024, dtype=torch.uint8, device="cuda") + for _ in range(2 * args.num_tests): + flush.zero_() + cuda.synchronize() + start = cuda.Event(enable_timing=True) + end = cuda.Event(enable_timing=True) + start.record() + alternate() + end.record() + end.synchronize() + times_ms.append(start.elapsed_time(end)) + samples: tuple[list[float], list[float]] = ([], []) + for mode, milliseconds in zip(order[-len(times_ms) :], times_ms, strict=True): + if not math.isfinite(milliseconds) or milliseconds <= 0: + raise RuntimeError(f"Invalid GPU event sample: {milliseconds}") + samples[mode].append(milliseconds * 1000) + stats = {} + for label, values in zip(("native", "fused"), samples, strict=True): + if len(values) != args.num_tests: + raise RuntimeError("GPU timing harness returned an unexpected sample count") + p50 = _percentile(values, 0.5) + stats[label] = {"median_us": p50, "p50_us": p50, "p90_us": _percentile(values, 0.9), "samples": len(values)} + stats["speedup_p50"] = stats["native"]["p50_us"] / stats["fused"]["p50_us"] + return stats diff --git a/benchmarks/kernels/benchmark_magi2_mhc.py b/benchmarks/kernels/benchmark_magi2_mhc.py new file mode 100644 index 000000000..c141f50b0 --- /dev/null +++ b/benchmarks/kernels/benchmark_magi2_mhc.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Compare native and fused MAGI-2 mHC post-processing on CUDA or MUSA. + +Run from the checkout root with PyTorch and the selected backend dependencies. +torchada is loaded automatically only on MUSA:: + + python -m benchmarks.kernels.benchmark_magi2_mhc \ + --tokens 1 4 16 64 129 4096 --hidden-size 3072 --num-tests 30 + +The released config supplies hidden=3072, four streams and 20 Sinkhorn steps. +Token counts are synthetic; supply per-rank sizes from the target trace when +comparing production shapes. FP32 post/residual logits are strided views of the +same [tokens, 24] projection, as in MHCHandler.compute_logits. Mix inputs use +the selected dtype. Both mix implementations consume identical coefficients. + +The three comparisons cover coefficient preparation, stream mixing, and their +combined pipeline (including intermediate allocations, excluding projection, +normalization and apply_pre). This is an eager operator diagnostic, without +compile or graph capture. It does not measure full model performance. +No CUDA/MUSA device (or no PyTorch) produces a JSON skip record and exit status zero. +""" + +import argparse +import math +from functools import partial +from typing import Any + +from benchmarks.kernels._magi2_gpu_bench import ( + add_arguments, + check_pair, + emit, + load_musa, + measure_pair, + report_environment, +) + + +def benchmark_shape(torch: Any, mhc: Any, args: argparse.Namespace, tokens: int) -> None: + dtype = getattr(torch, args.dtype) + device = torch.device("cuda") + packed = torch.randn(tokens, 24, device=device, dtype=torch.float32) + # Preserve the token stride of the model's shared projection allocation. + post_inputs = ( + packed[:, 4:8], + packed[:, 8:].view(tokens, 4, 4), + torch.tensor([0.8], device=device), + torch.randn(4, device=device), + torch.tensor([1.1], device=device), + torch.randn(4, 4, device=device), + ) + streams = torch.randn(tokens, 4, args.hidden_size, device=device, dtype=dtype) + branch = torch.randn(tokens, args.hidden_size, device=device, dtype=dtype) + post_op, mix_op = mhc.MHCPostResidual(), mhc.MHCMix() + scale = 1 / math.sqrt(4 * args.hidden_size) + kwargs = dict(scale=scale, iterations=args.sinkhorn_iterations, epsilon=1e-12, out_dtype=dtype) + prepare_reference = partial(post_op.forward_native, *post_inputs, **kwargs) + prepare_fused = partial(post_op, *post_inputs, **kwargs) + coefficients = prepare_reference() + mix_reference = partial(mix_op.forward_native, streams, branch, *coefficients) + mix_fused = partial(mix_op, streams, branch, *coefficients) + post_kernel, mix_kernel = "_mhc_post_residual_kernel", "_mhc_mix_kernel" + cases = { + "post_residual": (prepare_reference, prepare_fused, (post_kernel,)), + "mix": (mix_reference, mix_fused, (mix_kernel,)), + "pipeline": ( + lambda: mix_op.forward_native(streams, branch, *prepare_reference()), + lambda: mix_op(streams, branch, *prepare_fused()), + (post_kernel, mix_kernel), + ), + } + rtol = {torch.bfloat16: 1.6e-2, torch.float16: 2e-3, torch.float32: 1e-5}[dtype] + for name, (reference, fused, kernels) in cases.items(): + if args.op != "all" and args.op != name: + continue + # BF16 MUSA custom ops can quantize values near zero to the adjacent + # representable value; retain the strict relative gate while allowing + # one BF16 ulp in the absolute comparison. + atol = ( + (1e-2 if dtype == torch.bfloat16 else 1e-6) + if name == "post_residual" + else (1e-2 if dtype == torch.bfloat16 else 1e-5) + ) + parity = check_pair(torch, reference, fused, mhc, kernels, rtol=rtol, atol=atol) + timing = measure_pair(torch, reference, fused, args) + # Logical minimum traffic: input reads plus output writes. Native + # intermediates and repeat parameter loads can add actual DRAM traffic. + element_bytes = streams.element_size() + prepare_bytes = tokens * 20 * (4 + element_bytes) + 22 * 4 + mix_bytes = tokens * (9 * args.hidden_size + 20) * element_bytes + minimum_bytes = {"post_residual": prepare_bytes, "mix": mix_bytes, "pipeline": prepare_bytes + mix_bytes}[name] + # Sigmoid/exp costs and max reductions are not represented as FLOPs. + prepare_flops = tokens * (72 + 64 * args.sinkhorn_iterations) + mix_flops = tokens * args.hidden_size * 36 + scalar_flops = {"post_residual": prepare_flops, "mix": mix_flops, "pipeline": prepare_flops + mix_flops}[name] + emit( + status="pass", + op=name, + tokens=tokens, + hidden_size=args.hidden_size, + num_streams=4, + dtype=args.dtype, + logits_dtype="float32", + packed_shape=list(packed.shape), + post_stride=list(post_inputs[0].stride()), + residual_stride=list(post_inputs[1].stride()), + scale=scale, + sinkhorn_iterations=args.sinkhorn_iterations, + epsilon=1e-12, + minimum_io_bytes=minimum_bytes, + scalar_flops_excluding_exp_and_max=scalar_flops, + effective_fused_gbps=minimum_bytes / timing["fused"]["p50_us"] / 1000, + **parity, + **timing, + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + add_arguments(parser) + parser.add_argument("--hidden-size", type=int, default=3072) + parser.add_argument("--sinkhorn-iterations", type=int, default=20) + parser.add_argument("--op", choices=["all", "post_residual", "mix", "pipeline"], default="all") + args = parser.parse_args() + if args.hidden_size <= 0 or not 0 <= args.sinkhorn_iterations <= 64: + parser.error("hidden-size must be positive; sinkhorn-iterations must be in [0, 64]") + torch = load_musa(args, parser) + if torch is None: + return + cuda = getattr(torch, "cuda") + from vllm_omni.diffusion.layers import mhc + + report_environment(torch, args, mhc) + with torch.inference_mode(): + for tokens in args.tokens: + benchmark_shape(torch, mhc, args, tokens) + cuda.empty_cache() + + +if __name__ == "__main__": + main() diff --git a/tests/diffusion/layers/test_mhc.py b/tests/diffusion/layers/test_mhc.py new file mode 100644 index 000000000..135868dc7 --- /dev/null +++ b/tests/diffusion/layers/test_mhc.py @@ -0,0 +1,374 @@ +# 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) diff --git a/vllm_omni/diffusion/layers/mhc.py b/vllm_omni/diffusion/layers/mhc.py new file mode 100644 index 000000000..3fa6981b1 --- /dev/null +++ b/vllm_omni/diffusion/layers/mhc.py @@ -0,0 +1,279 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# Copyright (c) 2026 SandAI. All Rights Reserved. + +"""Native and eager Triton mHC post-processing, without projection or tensor caches.""" + +import threading + +import torch +from vllm.logger import init_logger +from vllm.triton_utils import HAS_TRITON, tl, triton + +from vllm_omni.diffusion.layers.custom_op import CustomOp + +_FAILED_MHC_KERNELS: set[tuple[str, int, str]] = set() +_WARNED_MHC_KERNELS: set[tuple[str, int, str]] = set() +_MHC_WARNING_LOCK = threading.Lock() +logger = init_logger(__name__) + + +def _mhc_kernel_key(tensor: torch.Tensor, name: str) -> tuple[str, int, str]: + return (tensor.device.type, tensor.device.index or 0, name) + + +def _is_mhc_oom_error(error: Exception) -> bool: + message = str(error).lower() + error_name = type(error).__name__.lower() + return isinstance(error, MemoryError) or "out of memory" in message or "outofmemoryerror" in error_name + + +def _cache_mhc_kernel_failure(failure_key: tuple[str, int, str], error: Exception) -> None: + _FAILED_MHC_KERNELS.add(failure_key) + with _MHC_WARNING_LOCK: + should_warn = failure_key not in _WARNED_MHC_KERNELS + _WARNED_MHC_KERNELS.add(failure_key) + if should_warn: + logger.warning( + "mHC fused kernel failed for key=%s; falling back to native: %s", + failure_key, + repr(error), + ) + + +def sinkhorn_knopp(matrix_logits: torch.Tensor, iterations: int, epsilon: float) -> torch.Tensor: + matrix = torch.exp(matrix_logits - matrix_logits.amax(dim=(-2, -1), keepdim=True)) + for _ in range(iterations): + matrix = matrix / (matrix.sum(dim=-2, keepdim=True) + epsilon) + matrix = matrix / (matrix.sum(dim=-1, keepdim=True) + epsilon) + return matrix + + +@triton.jit +def _mhc_post_residual_kernel( + post_ptr, + residual_ptr, + alpha_post_ptr, + bias_post_ptr, + alpha_residual_ptr, + bias_residual_ptr, + post_out_ptr, + residual_out_ptr, + post_stride, + residual_stride, + scale, + epsilon, + iterations: tl.constexpr, +): + token = tl.program_id(0).to(tl.int64) + index = tl.arange(0, 4) + post = tl.load(post_ptr + token * post_stride + index) + ap = tl.load(alpha_post_ptr) + bp = tl.load(bias_post_ptr + index) + post = 2.0 * tl.sigmoid(ap * scale * post + bp) + offsets = index[:, None] * 4 + index[None, :] + matrix = tl.load(residual_ptr + token * residual_stride + offsets) + ar = tl.load(alpha_residual_ptr) + br = tl.load(bias_residual_ptr + offsets) + matrix = ar * scale * matrix + br + # Native amax propagates any NaN to the whole token's normalization. + has_nan = tl.sum(tl.sum((matrix != matrix).to(tl.int32), axis=0), axis=0) != 0 + maximum = tl.max(tl.max(matrix, axis=0), axis=0) + maximum = tl.where(has_nan, float("nan"), maximum) + matrix = tl.exp(matrix - maximum) + for _ in tl.static_range(iterations): + matrix = matrix / (tl.sum(matrix, axis=0)[None, :] + epsilon) + matrix = matrix / (tl.sum(matrix, axis=1)[:, None] + epsilon) + tl.store(post_out_ptr + token * 4 + index, post) + tl.store(residual_out_ptr + token * 16 + offsets, matrix) + + +@triton.jit +def _mhc_mix_kernel(streams_ptr, branch_ptr, post_ptr, matrix_ptr, out_ptr, hidden, block_size: tl.constexpr): + token = tl.program_id(0).to(tl.int64) + channels = tl.program_id(1).to(tl.int64) * block_size + tl.arange(0, block_size) + valid = channels < hidden + stream_ids = tl.arange(0, 4) + mixed = tl.zeros((4, block_size), tl.float32) + for j in tl.static_range(4): + x = tl.load(streams_ptr + (token * 4 + j) * hidden + channels, mask=valid, other=0).to(tl.float32) + weight = tl.load(matrix_ptr + token * 16 + stream_ids * 4 + j).to(tl.float32) + mixed += weight[:, None] * x[None, :] + branch = tl.load(branch_ptr + token * hidden + channels, mask=valid, other=0).to(tl.float32) + post = tl.load(post_ptr + token * 4 + stream_ids).to(tl.float32) + branch = post[:, None] * branch[None, :] + # Preserve native einsum materialization before the final addition. + mixed = mixed.to(out_ptr.dtype.element_ty).to(tl.float32) + branch = branch.to(out_ptr.dtype.element_ty).to(tl.float32) + tl.store( + out_ptr + (token * 4 + stream_ids[:, None]) * hidden + channels[None, :], mixed + branch, mask=valid[None, :] + ) + + +def _native_required(tensors: tuple[torch.Tensor, ...]) -> bool: + return ( + torch.compiler.is_compiling() + or tensors[0].device.type not in ("cuda", "musa") + or (torch.is_grad_enabled() and any(t.requires_grad for t in tensors)) + ) + + +class MHCPostResidual(CustomOp): + """Prepare post coefficients and the residual matrix with FP32 arithmetic.""" + + @staticmethod + def forward_native( + post_logits: torch.Tensor, + residual_logits: torch.Tensor, + alpha_post: torch.Tensor, + bias_post: torch.Tensor, + alpha_residual: torch.Tensor, + bias_residual: torch.Tensor, + *, + scale: float, + iterations: int = 20, + epsilon: float = 1e-12, + out_dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor]: + post = 2.0 * torch.sigmoid(alpha_post * scale * post_logits + bias_post.unsqueeze(0)) + residual = sinkhorn_knopp( + alpha_residual * scale * residual_logits.float() + bias_residual.unsqueeze(0).float(), + iterations, + epsilon, + ) + return post.to(out_dtype), residual.to(out_dtype) + + def forward_cuda( + self, + post_logits: torch.Tensor, + residual_logits: torch.Tensor, + alpha_post: torch.Tensor, + bias_post: torch.Tensor, + alpha_residual: torch.Tensor, + bias_residual: torch.Tensor, + *, + scale: float, + iterations: int = 20, + epsilon: float = 1e-12, + out_dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor]: + tensors = (post_logits, residual_logits, alpha_post, bias_post, alpha_residual, bias_residual) + if _native_required(tensors): + return self.forward_native( + *tensors, scale=scale, iterations=iterations, epsilon=epsilon, out_dtype=out_dtype + ) + failure_key = _mhc_kernel_key(post_logits, "post_residual") + if ( + not HAS_TRITON + or failure_key in _FAILED_MHC_KERNELS + or not ( + post_logits.ndim == 2 + and post_logits.shape[1] == 4 + and residual_logits.shape == (post_logits.shape[0], 4, 4) + and post_logits.stride(1) == 1 + and residual_logits.stride()[1:] == (4, 1) + and alpha_post.numel() == 1 + and alpha_residual.numel() == 1 + and alpha_post.ndim <= 1 + and alpha_residual.ndim <= 1 + and bias_post.shape == (4,) + and bias_residual.shape == (4, 4) + and all(t.is_contiguous() for t in tensors[2:]) + and all(t.dtype == torch.float32 and t.device == post_logits.device for t in tensors) + and out_dtype in (torch.float32, torch.bfloat16, torch.float16) + and isinstance(scale, (int, float)) + and isinstance(epsilon, (int, float)) + and isinstance(iterations, int) + and 0 <= iterations <= 64 + ) + ): + return self.forward_native( + *tensors, scale=scale, iterations=iterations, epsilon=epsilon, out_dtype=out_dtype + ) + tokens = post_logits.shape[0] + post_out = torch.empty((tokens, 4), device=post_logits.device, dtype=out_dtype) + residual_out = torch.empty((tokens, 4, 4), device=post_logits.device, dtype=out_dtype) + if tokens: + try: + _mhc_post_residual_kernel[(tokens,)]( + *tensors, + post_out, + residual_out, + post_logits.stride(0), + residual_logits.stride(0), + scale, + epsilon, + iterations=iterations, + num_warps=1, + enable_fp_fusion=False, + ) + except Exception as error: + if _is_mhc_oom_error(error): + raise + _cache_mhc_kernel_failure(failure_key, error) + return self.forward_native( + *tensors, scale=scale, iterations=iterations, epsilon=epsilon, out_dtype=out_dtype + ) + return post_out, residual_out + + forward_npu = forward_native + + +class MHCMix(CustomOp): + """Mix four streams and add a separately rounded branch term.""" + + @staticmethod + def forward_native( + streams: torch.Tensor, + branch_output: torch.Tensor, + post_coefficients: torch.Tensor, + residual_matrix: torch.Tensor, + ) -> torch.Tensor: + branch = torch.einsum("tn,tc->tnc", post_coefficients, branch_output) + mixed = torch.einsum("tij,tjc->tic", residual_matrix, streams) + return mixed + branch + + def forward_cuda( + self, + streams: torch.Tensor, + branch_output: torch.Tensor, + post_coefficients: torch.Tensor, + residual_matrix: torch.Tensor, + ) -> torch.Tensor: + tensors = (streams, branch_output, post_coefficients, residual_matrix) + if _native_required(tensors): + return self.forward_native(*tensors) + failure_key = _mhc_kernel_key(streams, "mix") + if ( + not HAS_TRITON + or failure_key in _FAILED_MHC_KERNELS + or not ( + streams.ndim == 3 + and streams.shape[1] == 4 + and branch_output.shape == (streams.shape[0], streams.shape[2]) + and post_coefficients.shape == streams.shape[:2] + and residual_matrix.shape == (streams.shape[0], 4, 4) + and all(t.is_contiguous() and t.device == streams.device and t.dtype == streams.dtype for t in tensors) + and streams.dtype in (torch.float32, torch.bfloat16, torch.float16) + ) + ): + return self.forward_native(*tensors) + output = torch.empty_like(streams) + if output.numel(): + try: + _mhc_mix_kernel[(streams.shape[0], triton.cdiv(streams.shape[2], 256))]( + *tensors, + output, + streams.shape[2], + block_size=256, + num_warps=4, + enable_fp_fusion=False, + ) + except Exception as error: + if _is_mhc_oom_error(error): + raise + _cache_mhc_kernel_failure(failure_key, error) + return self.forward_native(*tensors) + return output + + forward_npu = forward_native diff --git a/vllm_omni/diffusion/models/magi2/layers.py b/vllm_omni/diffusion/models/magi2/layers.py index 3fe77ab69..a7de3b22a 100644 --- a/vllm_omni/diffusion/models/magi2/layers.py +++ b/vllm_omni/diffusion/models/magi2/layers.py @@ -19,8 +19,13 @@ import torch.distributed as dist import torch.nn as nn import torch.nn.functional as F +from vllm_omni.diffusion.layers.mhc import MHCMix, MHCPostResidual, sinkhorn_knopp + from .parallel import Magi2ParallelGroup, get_magi2_tp_group +_mhc_post_residual = MHCPostResidual() +_mhc_mix = MHCMix() + def swiglu7( x: torch.Tensor, @@ -329,14 +334,6 @@ class ElementWiseFourierEmbed(nn.Module): MHCTensorTuple = tuple[torch.Tensor, torch.Tensor, torch.Tensor] -def sinkhorn_knopp(matrix_logits: torch.Tensor, iterations: int, epsilon: float) -> torch.Tensor: - matrix = torch.exp(matrix_logits - matrix_logits.amax(dim=(-2, -1), keepdim=True)) - for _ in range(iterations): - matrix = matrix / (matrix.sum(dim=-2, keepdim=True) + epsilon) - matrix = matrix / (matrix.sum(dim=-1, keepdim=True) + epsilon) - return matrix - - class MHCHandler: """Exact four-stream manifold-constrained hyper-connection math.""" @@ -397,13 +394,19 @@ class MHCHandler: ) -> tuple[torch.Tensor, torch.Tensor]: alpha_post, bias_post, post_logits = post alpha_residual, bias_residual, residual_logits = residual - post_coefficients = 2.0 * torch.sigmoid(alpha_post * self.matmul_scale * post_logits + bias_post.unsqueeze(0)) - residual_matrix = sinkhorn_knopp( - alpha_residual * self.matmul_scale * residual_logits.float() + bias_residual.unsqueeze(0).float(), - self.sinkhorn_iterations, - self.sinkhorn_epsilon, + prepare = _mhc_post_residual if not torch.compiler.is_compiling() else _mhc_post_residual.forward_native + return prepare( + post_logits, + residual_logits, + alpha_post, + bias_post, + alpha_residual, + bias_residual, + scale=self.matmul_scale, + iterations=self.sinkhorn_iterations, + epsilon=self.sinkhorn_epsilon, + out_dtype=out_dtype, ) - return post_coefficients.to(out_dtype), residual_matrix.to(out_dtype) def hyper_connect( self, @@ -415,9 +418,8 @@ class MHCHandler: self._check_multi(residual_streams) if branch_output.ndim != 2 or branch_output.shape[-1] != self.hidden_size: raise ValueError("invalid mHC branch-output shape") - branch = torch.einsum("tn,tc->tnc", post_coefficients, branch_output) - mixed = torch.einsum("tij,tjc->tic", residual_matrix, residual_streams) - return mixed + branch + mix = _mhc_mix if not torch.compiler.is_compiling() else _mhc_mix.forward_native + return mix(residual_streams, branch_output, post_coefficients, residual_matrix) def _check_multi(self, tensor: torch.Tensor) -> None: expected = (self.num_streams, self.hidden_size)