| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import torch |
| 4 | |
| 5 | from vllm.triton_utils import tl, triton |
| 6 | |
| 7 | |
| 8 | @triton.jit |
| 9 | def _rms_norm(x, w, eps, HIDDEN_SIZE: tl.constexpr): |
| 10 | x = x.to(tl.float32) |
| 11 | mean_sq = tl.sum(x * x, axis=0) / HIDDEN_SIZE |
| 12 | rrms = tl.rsqrt(mean_sq + eps) |
| 13 | w = w.to(tl.float32) |
| 14 | return (x * rrms) * w |
| 15 | |
| 16 | |
| 17 | @triton.jit |
| 18 | def _fused_eh_norm_kernel( |
| 19 | pos_ptr, |
| 20 | embeds_ptr, |
| 21 | embeds_stride, |
| 22 | prev_ptr, |
| 23 | prev_stride, |
| 24 | enorm_w_ptr, |
| 25 | hnorm_w_ptr, |
| 26 | eps, |
| 27 | out_ptr, |
| 28 | out_stride, |
| 29 | H: tl.constexpr, |
| 30 | BLOCK: tl.constexpr, |
| 31 | ): |
| 32 | """MTP input fusion: zero embeds at position 0, RMSNorm(embeds) with enorm |
| 33 | and RMSNorm(prev_hidden) with hnorm, written side-by-side into ``out`` |
| 34 | ([N, 2H]) ready for the eh_proj GEMM. Replaces where + 2x RMSNorm + cat.""" |
| 35 | tok = tl.program_id(0) |
| 36 | off = tl.arange(0, BLOCK) |
| 37 | mask = off < H |
| 38 | |
| 39 | pos = tl.load(pos_ptr + tok) |
| 40 | e = tl.load(embeds_ptr + tok * embeds_stride + off, mask=mask, other=0.0) |
| 41 | e = tl.where(pos == 0, 0.0, e.to(tl.float32)) |
| 42 | ew = tl.load(enorm_w_ptr + off, mask=mask) |
| 43 | e_normed = _rms_norm(e, ew, eps, H) |
| 44 | tl.store(out_ptr + tok * out_stride + off, e_normed, mask=mask) |
| 45 | |
| 46 | p = tl.load(prev_ptr + tok * prev_stride + off, mask=mask, other=0.0) |
| 47 | hw = tl.load(hnorm_w_ptr + off, mask=mask) |
| 48 | p_normed = _rms_norm(p, hw, eps, H) |
| 49 | tl.store(out_ptr + tok * out_stride + H + off, p_normed, mask=mask) |
| 50 | |
| 51 | |
| 52 | def fused_eh_norm( |
| 53 | positions: torch.Tensor, |
| 54 | inputs_embeds: torch.Tensor, |
| 55 | previous_hidden: torch.Tensor, |
| 56 | enorm_w: torch.Tensor, |
| 57 | hnorm_w: torch.Tensor, |
| 58 | eps: float, |
| 59 | ) -> torch.Tensor: |
| 60 | """Returns cat([enorm(masked embeds), hnorm(prev_hidden)]) -> [N, 2H].""" |
| 61 | n, h = inputs_embeds.shape |
| 62 | out = torch.empty(n, 2 * h, dtype=inputs_embeds.dtype, device=inputs_embeds.device) |
| 63 | _fused_eh_norm_kernel[(n,)]( |
| 64 | positions, |
| 65 | inputs_embeds, |
| 66 | inputs_embeds.stride(0), |
| 67 | previous_hidden, |
| 68 | previous_hidden.stride(0), |
| 69 | enorm_w, |
| 70 | hnorm_w, |
| 71 | eps, |
| 72 | out, |
| 73 | out.stride(0), |
| 74 | h, |
| 75 | triton.next_power_of_2(h), |
| 76 | ) |
| 77 | return out |