From e75278dccc36ae894159d519f18c3b3db578a8bd Mon Sep 17 00:00:00 2001 From: R0CKSTAR Date: Sat, 26 Sep 2026 22:53:04 +0800 Subject: [PATCH] feat(magi2): add fused SwiGLU7 activation (#7252) Signed-off-by: Xiaodong Ye --- benchmarks/kernels/_magi2_musa_bench.py | 194 ++++++++++++++++++ benchmarks/kernels/benchmark_magi2_swiglu7.py | 87 ++++++++ tests/diffusion/layers/test_swiglu7.py | 194 ++++++++++++++++++ vllm_omni/diffusion/layers/swiglu7.py | 84 ++++++++ vllm_omni/diffusion/models/magi2/layers.py | 13 +- 5 files changed, 566 insertions(+), 6 deletions(-) create mode 100644 benchmarks/kernels/_magi2_musa_bench.py create mode 100644 benchmarks/kernels/benchmark_magi2_swiglu7.py create mode 100644 tests/diffusion/layers/test_swiglu7.py create mode 100644 vllm_omni/diffusion/layers/swiglu7.py diff --git a/benchmarks/kernels/_magi2_musa_bench.py b/benchmarks/kernels/_magi2_musa_bench.py new file mode 100644 index 000000000..cdfde092e --- /dev/null +++ b/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 diff --git a/benchmarks/kernels/benchmark_magi2_swiglu7.py b/benchmarks/kernels/benchmark_magi2_swiglu7.py new file mode 100644 index 000000000..0aade09ca --- /dev/null +++ b/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() diff --git a/tests/diffusion/layers/test_swiglu7.py b/tests/diffusion/layers/test_swiglu7.py new file mode 100644 index 000000000..bbb686397 --- /dev/null +++ b/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) diff --git a/vllm_omni/diffusion/layers/swiglu7.py b/vllm_omni/diffusion/layers/swiglu7.py new file mode 100644 index 000000000..8d9b69dcb --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/models/magi2/layers.py b/vllm_omni/diffusion/models/magi2/layers.py index a7de3b22a..8fcf39b2d 100644 --- a/vllm_omni/diffusion/models/magi2/layers.py +++ b/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: