| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Inkling dense SwiGLU MLP (also used as the MoE shared expert). |
| 4 | |
| 5 | The checkpoint stores the gate/up projection as a single fused, *interleaved* |
| 6 | weight (``[gate0, up0, gate1, up1, ...]``), so we use a plain |
| 7 | ``ColumnParallelLinear`` whose contiguous row-sharding keeps each gate/up pair |
| 8 | together, and an interleaved SwiGLU activation. |
| 9 | """ |
| 10 | |
| 11 | from __future__ import annotations |
| 12 | |
| 13 | import torch |
| 14 | from torch import nn |
| 15 | |
| 16 | from vllm.model_executor.layers.linear import ( |
| 17 | ColumnParallelLinear, |
| 18 | RowParallelLinear, |
| 19 | ) |
| 20 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 21 | |
| 22 | |
| 23 | class InklingDenseMLP(nn.Module): |
| 24 | def __init__( |
| 25 | self, |
| 26 | hidden_size: int, |
| 27 | intermediate_size: int, |
| 28 | *, |
| 29 | use_global_scale: bool = False, |
| 30 | quant_config: QuantizationConfig | None = None, |
| 31 | prefix: str = "", |
| 32 | ) -> None: |
| 33 | super().__init__() |
| 34 | self.gate_up_proj = ColumnParallelLinear( |
| 35 | hidden_size, |
| 36 | 2 * intermediate_size, |
| 37 | bias=False, |
| 38 | quant_config=quant_config, |
| 39 | prefix=f"{prefix}.gate_up_proj", |
| 40 | ) |
| 41 | self.down_proj = RowParallelLinear( |
| 42 | intermediate_size, |
| 43 | hidden_size, |
| 44 | bias=False, |
| 45 | quant_config=quant_config, |
| 46 | reduce_results=False, |
| 47 | prefix=f"{prefix}.down_proj", |
| 48 | ) |
| 49 | if use_global_scale: |
| 50 | self.global_scale = nn.Parameter(torch.empty(1), requires_grad=False) |
| 51 | else: |
| 52 | self.global_scale = None |
| 53 | |
| 54 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 55 | from .ops import silu_and_mul_triton |
| 56 | |
| 57 | gate_up, _ = self.gate_up_proj(x) |
| 58 | x = silu_and_mul_triton(gate_up) |
| 59 | x, _ = self.down_proj(x) |
| 60 | if self.global_scale is not None: |
| 61 | x.mul_(self.global_scale) |
| 62 | # TP-partial output: the layer's reduce-scatter fallback consumes it. |
| 63 | return x |