feat(magi2): add fused SwiGLU7 activation (#7252)

Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>

benchmarks/kernels/_magi2_musa_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_swiglu7.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Compare native and fused MAGI-2 SwiGLU7 operators on CUDA or MUSA.
+
+Run from the checkout root in an environment with PyTorch and the selected
+backend dependencies.  torchada is loaded automatically only on MUSA::
+
+    python -m benchmarks.kernels.benchmark_magi2_swiglu7 \
+        --tokens 1 4 16 64 129 4096 --intermediate-size 1280 --num-tests 30
+
+The released MAGI-2 config supplies the 1280 intermediate width. Token counts
+are a synthetic sweep; pass per-rank sizes from the target trace for production
+comparisons. Input is contiguous interleaved gate/up [tokens, 2 * intermediate].
+Both paths evaluate FP32 arithmetic with one cast to the requested output dtype.
+This measures eager operator latency, not compiled serving performance.
+No CUDA/MUSA device (or no PyTorch) produces a JSON skip record and exit status zero.
+"""
+
+import argparse
+from functools import partial
+
+from benchmarks.kernels._magi2_musa_bench import (
+    add_arguments,
+    check_pair,
+    emit,
+    load_musa,
+    measure_pair,
+    report_environment,
+)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    add_arguments(parser)
+    parser.add_argument("--intermediate-size", type=int, default=1280)
+    parser.add_argument("--out-dtype", choices=["bfloat16", "float16", "float32"], default=None)
+    args = parser.parse_args()
+    if args.intermediate_size <= 0:
+        parser.error("intermediate-size must be positive")
+    torch = load_musa(args, parser)
+    if torch is None:
+        return
+    cuda = getattr(torch, "cuda")
+    from vllm_omni.diffusion.layers import swiglu7
+
+    report_environment(torch, args, swiglu7)
+    dtype = getattr(torch, args.dtype)
+    out_dtype = getattr(torch, args.out_dtype or args.dtype)
+    device = torch.device("cuda")
+    rtol = {torch.bfloat16: 1.6e-2, torch.float16: 2e-3, torch.float32: 1e-5}[out_dtype]
+    atol = 1e-5 if out_dtype == torch.float32 else 2e-3
+    op = swiglu7.SwiGLU7()
+    with torch.inference_mode():
+        for tokens in args.tokens:
+            # Values outside both clamp boundaries exercise the released formula.
+            x = (torch.randn(tokens, 2 * args.intermediate_size, device=device) * 10).to(dtype)
+            reference = partial(op.forward_native, x, out_dtype=out_dtype)
+            fused = partial(op, x, out_dtype=out_dtype)
+            parity = check_pair(torch, reference, fused, swiglu7, ("_swiglu7_kernel",), rtol=rtol, atol=atol)
+            timing = measure_pair(torch, reference, fused, args)
+            output_elements = tokens * args.intermediate_size
+            minimum_io_bytes = (
+                x.numel() * x.element_size() + output_elements * torch.empty((), dtype=out_dtype).element_size()
+            )
+            emit(
+                status="pass",
+                op="swiglu7",
+                input_shape=list(x.shape),
+                input_stride=list(x.stride()),
+                dtype=args.dtype,
+                out_dtype=args.out_dtype or args.dtype,
+                alpha=1.702,
+                limit=7.0,
+                minimum_io_bytes=minimum_io_bytes,
+                # Six scalar arithmetic operations and one exp per output;
+                # clamps and transcendental costs are not modeled as FLOPs.
+                scalar_flops_excluding_exp_and_clamps=6 * output_elements,
+                effective_fused_gbps=minimum_io_bytes / timing["fused"]["p50_us"] / 1000,
+                **parity,
+                **timing,
+            )
+            del x, reference, fused
+            cuda.empty_cache()
+
+
+if __name__ == "__main__":
+    main()

tests/diffusion/layers/test_swiglu7.py

@@ -0,0 +1,194 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.layers import swiglu7 as kernel_module
+from vllm_omni.diffusion.layers.swiglu7 import SwiGLU7
+from vllm_omni.diffusion.models.magi2 import layers as magi2_layers
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.core_model]
+
+
+def reference(x, alpha=1.702, limit=7.0, out_dtype=None):
+    gate = x[..., ::2].float().clamp(max=limit)
+    up = x[..., 1::2].float().clamp(min=-limit, max=limit)
+    return (gate * torch.sigmoid(alpha * gate) * (up + 1)).to(x.dtype if out_dtype is None else out_dtype)
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
+@pytest.mark.parametrize("out_dtype", [None, torch.float32, torch.bfloat16])
+def test_native_dtype_and_rounding(dtype, out_dtype):
+    torch.manual_seed(71)
+    x = torch.randn(3, 14).to(dtype)
+    actual = SwiGLU7().forward_native(x, out_dtype=out_dtype)
+    expected = reference(x, out_dtype=out_dtype)
+    assert actual.dtype == expected.dtype
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("mode", ["noncontiguous", "empty", "rank3", "odd_broadcast", "odd_error", "custom"])
+def test_native_fallback_preserves_existing_inputs(mode):
+    op = SwiGLU7()
+    x = torch.randn(3, 12)
+    kwargs = {}
+    if mode == "noncontiguous":
+        x = x[:, ::2]
+    elif mode == "empty":
+        x = torch.empty(0, 12)
+    elif mode == "rank3":
+        x = torch.randn(2, 3, 12)
+    elif mode == "odd_broadcast":
+        x = torch.randn(3, 3)
+    elif mode == "odd_error":
+        with pytest.raises(RuntimeError):
+            op.forward_cuda(torch.randn(3, 7))
+        return
+    elif mode == "custom":
+        kwargs = dict(alpha=0.5, limit=2.0, out_dtype=torch.float16)
+    torch.testing.assert_close(op.forward_cuda(x, **kwargs), reference(x, **kwargs), rtol=0, atol=0)
+
+
+@pytest.mark.cpu
+def test_meta_output_shape():
+    x = torch.empty(2, 5, 18, device="meta", dtype=torch.bfloat16)
+    out = SwiGLU7().forward_cuda(x, out_dtype=torch.float32)
+    assert out.shape == (2, 5, 9) and out.dtype == torch.float32 and out.device.type == "meta"
+
+
+@pytest.mark.cpu
+def test_gradients_remain_native():
+    x = torch.randn(3, 12, requires_grad=True)
+    y = x.detach().clone().requires_grad_()
+    actual = SwiGLU7().forward_cuda(x)
+    expected = reference(y)
+    actual.sum().backward()
+    expected.sum().backward()
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    torch.testing.assert_close(x.grad, y.grad, rtol=0, atol=0)
+
+
+@pytest.mark.cpu
+def test_model_wiring_uses_fused_eager_path(monkeypatch):
+    x = torch.randn(3, 12)
+    calls = []
+
+    class Probe:
+        def __call__(self, x, alpha, limit, out_dtype):
+            calls.append("dispatch")
+            return reference(x, alpha, limit, out_dtype)
+
+        def forward_native(self, x, alpha, limit, out_dtype):
+            calls.append("native")
+            return reference(x, alpha, limit, out_dtype)
+
+    monkeypatch.setattr(magi2_layers, "_swiglu7_op", Probe())
+    torch.testing.assert_close(magi2_layers.swiglu7(x), reference(x))
+    assert calls == ["dispatch"]
+
+
+@pytest.mark.cpu
+def test_compile_keeps_native_pointwise_graph(monkeypatch):
+    torch._dynamo.reset()
+    targets: list[str] = []
+
+    def backend(gm, _inputs):
+        targets.extend(str(node.target) for node in gm.graph.nodes)
+        return gm.forward
+
+    compiled = torch.compile(magi2_layers.swiglu7, backend=backend, fullgraph=True)
+    x = torch.randn(3, 12)
+    torch.testing.assert_close(compiled(x), reference(x), rtol=0, atol=0)
+    assert any("sigmoid" in target for target in targets)
+    assert not any("triton" in target or "swiglu7_kernel" in target for target in targets)
+
+
+GPU_DEVICES = [pytest.param("cuda", marks=pytest.mark.cuda), pytest.param("musa", marks=pytest.mark.musa)]
+
+
+def require_device(name):
+    backend = getattr(torch, name, None)
+    if backend is None or not backend.is_available():
+        pytest.skip(f"requires {name}")
+    return torch.device(name)
+
+
+@pytest.mark.parametrize("device_name", GPU_DEVICES)
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
+@pytest.mark.parametrize("shape", [(129, 2560), (2, 514), (2, 3, 128), (0, 2560)])
+def test_eager_kernel_launch_and_numerics(monkeypatch, device_name, dtype, shape):
+    device = require_device(device_name)
+    torch.manual_seed(42)
+    x = (torch.randn(shape, dtype=torch.float32, device=device) * 10).to(dtype)
+    before = x.clone()
+    real_kernel = kernel_module._swiglu7_kernel
+    launches = []
+
+    class Spy:
+        def __getitem__(self, grid):
+            launches.append(grid)
+            return real_kernel[grid]
+
+    monkeypatch.setattr(kernel_module, "_swiglu7_kernel", Spy())
+    actual = SwiGLU7()(x)
+    getattr(torch, device_name).synchronize()
+    assert len(launches) == (1 if x.numel() else 0)
+    # Up to two output ULPs accommodate FP32 sigmoid implementation differences
+    # before the final BF16/FP16 rounding; FP32 keeps a tighter numerical gate.
+    rtol = {torch.bfloat16: 1.6e-2, torch.float16: 2e-3, torch.float32: 1e-5}[dtype]
+    atol = 1e-5 if dtype == torch.float32 else 2e-3
+    torch.testing.assert_close(actual, reference(x), rtol=rtol, atol=atol)
+    torch.testing.assert_close(x, before, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("device_name", GPU_DEVICES)
+def test_special_values_and_explicit_output_dtype(device_name):
+    device = require_device(device_name)
+    x = torch.tensor(
+        [
+            [
+                -100.0,
+                2.0,
+                -8.0,
+                1.0,
+                -7.0,
+                0.0,
+                0.0,
+                -1.0,
+                8.0,
+                100.0,
+                float("inf"),
+                -100.0,
+                float("-inf"),
+                0.0,
+                float("nan"),
+                2.0,
+                1.0,
+                float("nan"),
+            ]
+        ],
+        device=device,
+        dtype=torch.bfloat16,
+    )
+    actual = SwiGLU7()(x, out_dtype=torch.float32)
+    expected = reference(x, out_dtype=torch.float32)
+    getattr(torch, device_name).synchronize()
+    assert actual.dtype == torch.float32
+    torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5, equal_nan=True)
+
+
+@pytest.mark.parametrize("device_name", GPU_DEVICES)
+def test_gpu_compile_does_not_launch_standalone_kernel(monkeypatch, device_name):
+    device = require_device(device_name)
+
+    class ForbiddenKernel:
+        def __getitem__(self, _grid):
+            pytest.fail("compilation must use the visible native expression")
+
+    monkeypatch.setattr(kernel_module, "_swiglu7_kernel", ForbiddenKernel())
+    op = torch.compile(SwiGLU7(), backend="eager", fullgraph=True)
+    x = torch.randn(7, 128, device=device, dtype=torch.bfloat16)
+    torch.testing.assert_close(op(x), reference(x), rtol=0, atol=0)

vllm_omni/diffusion/layers/swiglu7.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# Copyright (c) 2026 SandAI. All Rights Reserved.
+
+"""Interleaved clamped SwiGLU7 with a native compile path.
+
+The Triton expression is adapted from SGLang revision
+25536524af6701518d0b8ec0efca8109c23b3706 (Apache-2.0).
+"""
+
+import torch
+from vllm.triton_utils import tl, triton
+
+from vllm_omni.diffusion.layers.custom_op import CustomOp
+
+
+@triton.jit
+def _swiglu7_kernel(x_ptr, out_ptr, numel, block_size: tl.constexpr):
+    offsets = tl.program_id(0).to(tl.int64) * block_size + tl.arange(0, block_size)
+    mask = offsets < numel
+    gate = tl.load(x_ptr + 2 * offsets, mask=mask, other=0).to(tl.float32)
+    linear = tl.load(x_ptr + 2 * offsets + 1, mask=mask, other=0).to(tl.float32)
+    # Comparisons preserve NaNs; gate is only upper-clamped, unlike linear.
+    gate = tl.where(gate > 7.0, 7.0, gate)
+    linear = tl.where(linear < -7.0, -7.0, linear)
+    linear = tl.where(linear > 7.0, 7.0, linear)
+    result = gate * tl.sigmoid(1.702 * gate) * (linear + 1.0)
+    tl.store(out_ptr + offsets, result, mask=mask)
+
+
+class SwiGLU7(CustomOp):
+    """SwiGLU7 for packed ``[..., gate_0, up_0, gate_1, up_1, ...]`` inputs.
+
+    Arithmetic is evaluated in FP32 and cast once to the requested output
+    dtype. Eager inference with contiguous inputs and the released constants
+    uses Triton; compilation, autograd and other layouts retain the native
+    expression so this pointwise operation does not become a fusion barrier.
+    """
+
+    @staticmethod
+    def forward_native(
+        x: torch.Tensor,
+        alpha: float = 1.702,
+        limit: float = 7.0,
+        out_dtype: torch.dtype | None = None,
+    ) -> torch.Tensor:
+        out_dtype = x.dtype if out_dtype is None else out_dtype
+        x = x.to(torch.float32)
+        gate, linear = x[..., ::2], x[..., 1::2]
+        gate = gate.clamp(max=limit)
+        linear = linear.clamp(min=-limit, max=limit)
+        return (gate * torch.sigmoid(alpha * gate) * (linear + 1.0)).to(out_dtype)
+
+    def forward_cuda(
+        self,
+        x: torch.Tensor,
+        alpha: float = 1.702,
+        limit: float = 7.0,
+        out_dtype: torch.dtype | None = None,
+    ) -> torch.Tensor:
+        if torch.compiler.is_compiling() or (torch.is_grad_enabled() and x.requires_grad):
+            return self.forward_native(x, alpha, limit, out_dtype)
+        supported_dtypes = (torch.float16, torch.bfloat16, torch.float32)
+        target_dtype = x.dtype if out_dtype is None else out_dtype
+        if (
+            x.device.type in ("cpu", "meta")
+            or x.ndim == 0
+            or x.shape[-1] % 2
+            or not x.is_contiguous()
+            or x.dtype not in supported_dtypes
+            or target_dtype not in supported_dtypes
+            or alpha != 1.702
+            or limit != 7.0
+        ):
+            return self.forward_native(x, alpha, limit, out_dtype)
+        output = torch.empty((*x.shape[:-1], x.shape[-1] // 2), device=x.device, dtype=target_dtype)
+        if output.numel():
+            _swiglu7_kernel[(triton.cdiv(output.numel(), 256),)](
+                x, output, output.numel(), block_size=256, num_warps=4, enable_fp_fusion=False
+            )
+        return output
+
+    # Platforms without a Triton implementation still have the native formula.
+    forward_npu = forward_native

vllm_omni/diffusion/models/magi2/layers.py

@@ -20,11 +20,13 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 from vllm_omni.diffusion.layers.mhc import MHCMix, MHCPostResidual, sinkhorn_knopp
+from vllm_omni.diffusion.layers.swiglu7 import SwiGLU7
 
 from .parallel import Magi2ParallelGroup, get_magi2_tp_group
 
 _mhc_post_residual = MHCPostResidual()
 _mhc_mix = MHCMix()
+_swiglu7_op = SwiGLU7()
 
 
 def swiglu7(
@@ -35,12 +37,11 @@ def swiglu7(
 ) -> torch.Tensor:
     """Released GPT-OSS-style clamped SwiGLU activation."""
 
-    out_dtype = x.dtype if out_dtype is None else out_dtype
-    x = x.to(torch.float32)
-    gate, linear = x[..., ::2], x[..., 1::2]
-    gate = gate.clamp(max=limit)
-    linear = linear.clamp(min=-limit, max=limit)
-    return (gate * torch.sigmoid(alpha * gate) * (linear + 1.0)).to(out_dtype)
+    # Keep the native expression visible to torch.compile while dispatching
+    # the fused CustomOp for eager inference.
+    if not torch.compiler.is_compiling():
+        return _swiglu7_op(x, alpha, limit, out_dtype)
+    return _swiglu7_op.forward_native(x, alpha, limit, out_dtype)
 
 
 class ModalityDispatcher: