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

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)

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

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)