| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Fused Kimi-K3 MTP input preparation.""" |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | from vllm.triton_utils import tl, triton |
| 8 | |
| 9 | |
| 10 | @triton.jit |
| 11 | def _rms_norm(x, weight, eps, hidden_size: tl.constexpr): |
| 12 | x = x.to(tl.float32) |
| 13 | variance = tl.sum(x * x, axis=0) / hidden_size |
| 14 | return x * tl.rsqrt(variance + eps) * weight.to(tl.float32) |
| 15 | |
| 16 | |
| 17 | @triton.jit |
| 18 | def _fused_mtp_input_kernel( |
| 19 | positions_ptr, |
| 20 | inputs_embeds_ptr, |
| 21 | previous_hidden_states_ptr, |
| 22 | enorm_weight_ptr, |
| 23 | hnorm_weight_ptr, |
| 24 | output_ptr, |
| 25 | eps, |
| 26 | inputs_embeds_stride, |
| 27 | previous_hidden_states_stride, |
| 28 | output_stride, |
| 29 | hidden_size: tl.constexpr, |
| 30 | block_size: tl.constexpr, |
| 31 | ): |
| 32 | token_idx = tl.program_id(0).to(tl.int64) |
| 33 | input_idx = tl.program_id(1) |
| 34 | offsets = tl.arange(0, block_size) |
| 35 | mask = offsets < hidden_size |
| 36 | |
| 37 | if input_idx == 0: |
| 38 | position = tl.load(positions_ptr + token_idx) |
| 39 | values = tl.load( |
| 40 | inputs_embeds_ptr + token_idx * inputs_embeds_stride + offsets, |
| 41 | mask=mask & (position != 0), |
| 42 | other=0.0, |
| 43 | ) |
| 44 | weight = tl.load(enorm_weight_ptr + offsets, mask=mask, other=0.0) |
| 45 | else: |
| 46 | values = tl.load( |
| 47 | previous_hidden_states_ptr |
| 48 | + token_idx * previous_hidden_states_stride |
| 49 | + offsets, |
| 50 | mask=mask, |
| 51 | other=0.0, |
| 52 | ) |
| 53 | weight = tl.load(hnorm_weight_ptr + offsets, mask=mask, other=0.0) |
| 54 | |
| 55 | output = _rms_norm(values, weight, eps, hidden_size) |
| 56 | tl.store( |
| 57 | output_ptr + token_idx * output_stride + input_idx * hidden_size + offsets, |
| 58 | output, |
| 59 | mask=mask, |
| 60 | ) |
| 61 | |
| 62 | |
| 63 | def fused_mtp_input( |
| 64 | positions: torch.Tensor, |
| 65 | inputs_embeds: torch.Tensor, |
| 66 | previous_hidden_states: torch.Tensor, |
| 67 | enorm_weight: torch.Tensor, |
| 68 | hnorm_weight: torch.Tensor, |
| 69 | eps: float, |
| 70 | ) -> torch.Tensor: |
| 71 | """Mask and normalize both MTP inputs into the projection layout.""" |
| 72 | num_tokens, hidden_size = inputs_embeds.shape |
| 73 | output = torch.empty( |
| 74 | num_tokens, |
| 75 | 2 * hidden_size, |
| 76 | dtype=inputs_embeds.dtype, |
| 77 | device=inputs_embeds.device, |
| 78 | ) |
| 79 | if num_tokens == 0: |
| 80 | return output |
| 81 | |
| 82 | _fused_mtp_input_kernel[(num_tokens, 2)]( |
| 83 | positions, |
| 84 | inputs_embeds, |
| 85 | previous_hidden_states, |
| 86 | enorm_weight, |
| 87 | hnorm_weight, |
| 88 | output, |
| 89 | eps, |
| 90 | inputs_embeds.stride(0), |
| 91 | previous_hidden_states.stride(0), |
| 92 | output.stride(0), |
| 93 | hidden_size, |
| 94 | triton.next_power_of_2(hidden_size), |
| 95 | ) |
| 96 | return output |