vllm-omni/benchmarks/kernels/benchmark_magi2_mhc.py

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