| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | import torch.nn.functional as F |
| 6 | from torch import Tensor |
| 7 | |
| 8 | from ..op import register_op |
| 9 | |
| 10 | |
| 11 | @register_op |
| 12 | def gelu_and_mul_sparse( |
| 13 | x: Tensor, std_multiplier: float, approximate: str = "none" |
| 14 | ) -> Tensor: |
| 15 | """Apply Gaussian sparsification, GELU, and gated multiplication.""" |
| 16 | d = x.shape[-1] // 2 |
| 17 | gate = x[..., :d] |
| 18 | # Statistics intentionally remain local to each tensor-parallel shard. |
| 19 | mean = torch.mean(gate, dim=-1, keepdim=True) |
| 20 | std = torch.std(gate, dim=-1, keepdim=True, unbiased=False) |
| 21 | sparse_gate = F.relu(gate - (mean + std * std_multiplier)) |
| 22 | return F.gelu(sparse_gate, approximate=approximate) * x[..., d:] |
| 23 | |
| 24 | |
| 25 | @gelu_and_mul_sparse.register_input_generator |
| 26 | def _gelu_and_mul_sparse_input_generator( |
| 27 | num_tokens: int, |
| 28 | hidden_size: int, |
| 29 | dtype: torch.dtype, |
| 30 | std_multiplier: float = 1.6448536269514722, |
| 31 | approximate: str = "tanh", |
| 32 | ) -> tuple: |
| 33 | x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype) |
| 34 | return x, std_multiplier, approximate |