| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """SwiGLU kernels for the Inkling MLP layers. |
| 4 | |
| 5 | ``silu_and_mul_triton``: SiLU-and-mul over the checkpoint's interleaved |
| 6 | fused gate/up layout (dense MLP). ``sink_silu_mul_epilogue``: the sink-expert |
| 7 | variant with the per-expert dequant scale and per-token gamma fused in. |
| 8 | """ |
| 9 | |
| 10 | import torch |
| 11 | |
| 12 | from vllm.triton_utils import tl, triton |
| 13 | |
| 14 | |
| 15 | @triton.jit(do_not_specialize=["M"]) |
| 16 | def _silu_and_mul_triton_kernel( |
| 17 | gateup_out_ptr, |
| 18 | down_inp_ptr, |
| 19 | M, |
| 20 | N: tl.constexpr, |
| 21 | GRID_SIZE: tl.constexpr, |
| 22 | NUM_STAGES: tl.constexpr, |
| 23 | BLOCK_SIZE_M: tl.constexpr, |
| 24 | BLOCK_SIZE_N: tl.constexpr, |
| 25 | EVEN_N: tl.constexpr, |
| 26 | INT64_INDEX: tl.constexpr, |
| 27 | ): |
| 28 | start_pid = tl.program_id(0) |
| 29 | if INT64_INDEX: |
| 30 | start_pid = start_pid.to(tl.int64) |
| 31 | M = M.to(tl.int64) |
| 32 | |
| 33 | NUM_BLOCKS_N: tl.constexpr = tl.cdiv(N, BLOCK_SIZE_N) |
| 34 | num_blocks_mn = tl.cdiv(M, BLOCK_SIZE_M) * NUM_BLOCKS_N |
| 35 | |
| 36 | for pid in tl.range(start_pid, num_blocks_mn, GRID_SIZE, num_stages=NUM_STAGES): |
| 37 | pid_m = pid // NUM_BLOCKS_N |
| 38 | pid_n = pid % NUM_BLOCKS_N |
| 39 | |
| 40 | offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) |
| 41 | offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) |
| 42 | mask_m = offs_m < M |
| 43 | mask_n = offs_n < N |
| 44 | |
| 45 | # Interleaved fused gate/up: [g0, u0, g1, u1, ...]. |
| 46 | mask_offs_2n = pid_n * BLOCK_SIZE_N + tl.arange(0, 2 * BLOCK_SIZE_N) // 2 |
| 47 | tl.static_assert(BLOCK_SIZE_N % 8 == 0, f"{BLOCK_SIZE_N=}") |
| 48 | mask_2n = mask_offs_2n < N |
| 49 | mask_2n = tl.max_constancy(mask_2n, [16]) |
| 50 | |
| 51 | offs_2n = pid_n * 2 * BLOCK_SIZE_N + tl.arange(0, 2 * BLOCK_SIZE_N) |
| 52 | offs_m2n = offs_m[:, None] * N * 2 + offs_2n[None, :] |
| 53 | |
| 54 | if EVEN_N or pid_n * BLOCK_SIZE_N + BLOCK_SIZE_N <= N: |
| 55 | gateup_out = tl.load( |
| 56 | gateup_out_ptr + offs_m2n, mask=mask_m[:, None], other=0.0 |
| 57 | ) |
| 58 | else: |
| 59 | mask_m2n = mask_m[:, None] & mask_2n[None, :] |
| 60 | gateup_out = tl.load(gateup_out_ptr + offs_m2n, mask=mask_m2n, other=0.0) |
| 61 | |
| 62 | gate_out, up_out = tl.split( |
| 63 | tl.reshape(gateup_out, (BLOCK_SIZE_M, BLOCK_SIZE_N, 2)) |
| 64 | ) |
| 65 | gate_out = gate_out.to(tl.float32) |
| 66 | up_out = up_out.to(tl.float32) |
| 67 | |
| 68 | down_inp = gate_out * tl.sigmoid(gate_out) * up_out |
| 69 | |
| 70 | mask_mn = mask_m[:, None] if EVEN_N else mask_m[:, None] & mask_n[None, :] |
| 71 | offs_mn = offs_m[:, None] * N + offs_n[None, :] |
| 72 | tl.store(down_inp_ptr + offs_mn, down_inp, mask=mask_mn) |
| 73 | |
| 74 | |
| 75 | def silu_and_mul_triton(gateup_output: torch.Tensor) -> torch.Tensor: |
| 76 | """SiLU-and-mul for the interleaved fused gate/up layout. |
| 77 | |
| 78 | Adapted from ``inkling_kernels.activation.silu_and_mul_fwd`` (without MXFP). |
| 79 | """ |
| 80 | assert gateup_output.is_contiguous(), ( |
| 81 | f"{gateup_output.shape=} {gateup_output.stride()=}" |
| 82 | ) |
| 83 | assert gateup_output.ndim == 2, f"{gateup_output.shape=}" |
| 84 | |
| 85 | M = gateup_output.shape[0] |
| 86 | hidden_size = gateup_output.shape[1] |
| 87 | assert hidden_size % 2 == 0, f"{hidden_size=}" |
| 88 | N = hidden_size // 2 |
| 89 | |
| 90 | down_input = torch.empty( |
| 91 | (M, N), device=gateup_output.device, dtype=gateup_output.dtype |
| 92 | ) |
| 93 | if M == 0: |
| 94 | return down_input |
| 95 | |
| 96 | BLOCK_SIZE_N = max(8, min(256, triton.next_power_of_2(N))) |
| 97 | if M <= 1: |
| 98 | BLOCK_SIZE_M = 4 |
| 99 | elif M <= 256: |
| 100 | BLOCK_SIZE_M = 2 |
| 101 | elif M < 4096: |
| 102 | BLOCK_SIZE_M = 4 |
| 103 | else: |
| 104 | BLOCK_SIZE_M = 16 |
| 105 | BLOCK_SIZE_N = max(8, min(128, triton.next_power_of_2(N))) |
| 106 | max_grid_size = triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N) |
| 107 | num_sms = torch.cuda.get_device_properties( |
| 108 | gateup_output.device |
| 109 | ).multi_processor_count |
| 110 | grid_size = min(num_sms * 4, max_grid_size) |
| 111 | |
| 112 | _silu_and_mul_triton_kernel[(grid_size,)]( |
| 113 | gateup_out_ptr=gateup_output, |
| 114 | down_inp_ptr=down_input, |
| 115 | M=M, |
| 116 | N=N, |
| 117 | GRID_SIZE=grid_size, |
| 118 | NUM_STAGES=1, |
| 119 | BLOCK_SIZE_M=BLOCK_SIZE_M, |
| 120 | BLOCK_SIZE_N=BLOCK_SIZE_N, |
| 121 | EVEN_N=N % BLOCK_SIZE_N == 0, |
| 122 | INT64_INDEX=gateup_output.nbytes >= 2**31, |
| 123 | num_warps=8, |
| 124 | ) |
| 125 | |
| 126 | return down_input |
| 127 | |
| 128 | |
| 129 | @triton.jit(do_not_specialize=["T"]) |
| 130 | def _sink_epilogue_kernel( |
| 131 | raw_ptr, # [T, S * 2F] gemm1 output, interleaved g/u pairs per expert block |
| 132 | alpha_ptr, # [S] fp32 per-expert pre-SiLU dequant scale |
| 133 | gamma_ptr, # [T, S] fp32 per-token sink weights (may be strided) |
| 134 | ratio_ptr, # [S] fp32 per-expert post-SiLU scale (gemm2 alpha ratio) |
| 135 | out_ptr, # [T, S * F] output |
| 136 | T, |
| 137 | stride_raw_0, |
| 138 | stride_gamma_0, |
| 139 | F: tl.constexpr, |
| 140 | S: tl.constexpr, |
| 141 | BLOCK_F: tl.constexpr, |
| 142 | ): |
| 143 | pid_t = tl.program_id(0).to(tl.int64) |
| 144 | pid_sf = tl.program_id(1) |
| 145 | if pid_t >= T: |
| 146 | return |
| 147 | s = pid_sf // (F // BLOCK_F) |
| 148 | offs_f = (pid_sf % (F // BLOCK_F)) * BLOCK_F + tl.arange(0, BLOCK_F) |
| 149 | |
| 150 | base = pid_t * stride_raw_0 + s * 2 * F |
| 151 | gate = tl.load(raw_ptr + base + 2 * offs_f).to(tl.float32) |
| 152 | up = tl.load(raw_ptr + base + 2 * offs_f + 1).to(tl.float32) |
| 153 | alpha = tl.load(alpha_ptr + s) |
| 154 | weight = tl.load(gamma_ptr + pid_t * stride_gamma_0 + s) * tl.load(ratio_ptr + s) |
| 155 | |
| 156 | gate *= alpha |
| 157 | up *= alpha |
| 158 | h = gate * tl.sigmoid(gate) * up * weight |
| 159 | tl.store(out_ptr + pid_t * (S * F) + s * F + offs_f, h) |
| 160 | |
| 161 | |
| 162 | def sink_silu_mul_epilogue( |
| 163 | raw: torch.Tensor, # [T, S * 2F] gemm1 output (interleaved gate/up rows) |
| 164 | alphas: torch.Tensor, # [S] fp32 |
| 165 | gammas: torch.Tensor, # [T, S] fp32 |
| 166 | ratios: torch.Tensor, # [S] fp32 |
| 167 | n_experts: int, |
| 168 | out_dtype: torch.dtype, |
| 169 | ) -> torch.Tensor: |
| 170 | """Fused sink-expert epilogue: silu(g * a_e) * (u * a_e) * (gamma * r_e). |
| 171 | |
| 172 | One kernel replaces the per-expert dequant column scale, the SwiGLU, and |
| 173 | the per-token gamma multiply between the two sink GEMMs. |
| 174 | """ |
| 175 | tokens = raw.shape[0] |
| 176 | f = raw.shape[1] // (2 * n_experts) |
| 177 | out = torch.empty((tokens, n_experts * f), device=raw.device, dtype=out_dtype) |
| 178 | if tokens == 0: |
| 179 | return out |
| 180 | # raw may be a column-slice of a padded GEMM output (rows strided). |
| 181 | assert raw.stride(1) == 1 and gammas.stride(1) == 1 |
| 182 | # Largest power-of-two divisor of f (f = 768 -> 256), capped at 512. |
| 183 | block_f = min(512, f & (-f)) |
| 184 | _sink_epilogue_kernel[(tokens, n_experts * (f // block_f))]( |
| 185 | raw, |
| 186 | alphas, |
| 187 | gammas, |
| 188 | ratios, |
| 189 | out, |
| 190 | tokens, |
| 191 | raw.stride(0), |
| 192 | gammas.stride(0), |
| 193 | F=f, |
| 194 | S=n_experts, |
| 195 | BLOCK_F=block_f, |
| 196 | ) |
| 197 | return out |