85 lines
3.3 KiB
Python
85 lines
3.3 KiB
Python
# 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
|