feat(magi2): add fused mHC post-processing (#7261)

Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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

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()