| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import torch |
| 4 | |
| 5 | |
| 6 | def mhc_pre_delayed_torch( |
| 7 | residual: torch.Tensor, |
| 8 | fn: torch.Tensor, |
| 9 | hc_scale: torch.Tensor, |
| 10 | hc_base: torch.Tensor, |
| 11 | rms_eps: float, |
| 12 | hc_pre_eps: float, |
| 13 | hc_sinkhorn_eps: float, |
| 14 | hc_post_mult_value: float, |
| 15 | sinkhorn_repeat: int, |
| 16 | pre_mix: torch.Tensor | None = None, |
| 17 | x: torch.Tensor | None = None, |
| 18 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 19 | """Reference for mHC pre using coefficients from the previous sublayer.""" |
| 20 | hc_mult = residual.shape[1] |
| 21 | x = (residual.flatten(1) if x is None else x).float() |
| 22 | mixes = (x @ fn.t()) * torch.rsqrt(x.square().mean(-1, keepdim=True) + rms_eps) |
| 23 | pre = ( |
| 24 | torch.sigmoid(mixes[:, :hc_mult] * hc_scale[0] + hc_base[:hc_mult]) + hc_pre_eps |
| 25 | ) |
| 26 | post = ( |
| 27 | torch.sigmoid( |
| 28 | mixes[:, hc_mult : 2 * hc_mult] * hc_scale[1] |
| 29 | + hc_base[hc_mult : 2 * hc_mult] |
| 30 | ) |
| 31 | * hc_post_mult_value |
| 32 | ) |
| 33 | comb = mixes[:, 2 * hc_mult :].view(-1, hc_mult, hc_mult) * hc_scale[2] |
| 34 | comb = comb + hc_base[2 * hc_mult :].view(1, hc_mult, hc_mult) |
| 35 | comb = torch.softmax(comb, dim=-1) + hc_sinkhorn_eps |
| 36 | comb = comb / (comb.sum(dim=-2, keepdim=True) + hc_sinkhorn_eps) |
| 37 | for _ in range(sinkhorn_repeat - 1): |
| 38 | comb = comb / (comb.sum(dim=-1, keepdim=True) + hc_sinkhorn_eps) |
| 39 | comb = comb / (comb.sum(dim=-2, keepdim=True) + hc_sinkhorn_eps) |
| 40 | layer_input = ( |
| 41 | residual[:, 0] |
| 42 | if pre_mix is None |
| 43 | else (pre_mix.unsqueeze(-1) * residual.float()).sum(dim=1).to(residual.dtype) |
| 44 | ) |
| 45 | return post.unsqueeze(-1), comb, layer_input, pre |
| 46 | |
| 47 | |
| 48 | def mhc_pre_torch( |
| 49 | residual: torch.Tensor, |
| 50 | fn: torch.Tensor, |
| 51 | hc_scale: torch.Tensor, |
| 52 | hc_base: torch.Tensor, |
| 53 | rms_eps: float, |
| 54 | hc_pre_eps: float, |
| 55 | hc_sinkhorn_eps: float, |
| 56 | hc_post_mult_value: float, |
| 57 | sinkhorn_repeat: int, |
| 58 | n_splits: int = 1, |
| 59 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 60 | """Forward pass for mHC pre block. |
| 61 | |
| 62 | Args: |
| 63 | residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16 |
| 64 | fn: shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32 |
| 65 | hc_scale: shape (3,), dtype torch.float32 |
| 66 | hc_base: shape (hc_mult3,), dtype torch.float32 |
| 67 | rms_eps: RMS normalization epsilon |
| 68 | hc_pre_eps: pre-mix epsilon |
| 69 | hc_sinkhorn_eps: sinkhorn epsilon |
| 70 | hc_post_mult_value: post-mix multiplier value |
| 71 | sinkhorn_repeat: number of sinkhorn iterations |
| 72 | n_splits: split-k factor; |
| 73 | |
| 74 | Returns: |
| 75 | post_mix: shape (..., hc_mult), dtype torch.float32 |
| 76 | comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32 |
| 77 | layer_input: shape (..., hidden_size), dtype torch.bfloat16 |
| 78 | |
| 79 | """ |
| 80 | # Validate shapes |
| 81 | assert residual.dtype == torch.bfloat16 |
| 82 | assert fn.dtype == torch.float32 |
| 83 | assert hc_scale.dtype == torch.float32 |
| 84 | assert hc_base.dtype == torch.float32 |
| 85 | |
| 86 | hc_mult = residual.shape[-2] |
| 87 | hidden_size = residual.shape[-1] |
| 88 | hc_mult2 = hc_mult * hc_mult |
| 89 | hc_mult3 = hc_mult * 2 + hc_mult2 |
| 90 | |
| 91 | hc_hidden_size = hc_mult * hidden_size |
| 92 | assert fn.shape[0] == hc_mult3 |
| 93 | assert fn.shape[1] == hc_hidden_size |
| 94 | assert hc_scale.shape == (3,) |
| 95 | assert hc_base.shape == (hc_mult3,) |
| 96 | |
| 97 | outer_shape = residual.shape[:-2] |
| 98 | |
| 99 | residual_flat = residual.view(-1, hc_mult, hidden_size) |
| 100 | num_tokens = residual_flat.shape[0] |
| 101 | fn_flat = fn |
| 102 | |
| 103 | x = residual_flat.view(num_tokens, hc_mult * hidden_size).to(torch.float32) |
| 104 | mixes = torch.matmul(x, fn_flat.t()) |
| 105 | sqrsum = x.square().sum(dim=-1, keepdim=True) |
| 106 | mixes = mixes * torch.rsqrt(sqrsum / (hc_mult * hidden_size) + rms_eps) |
| 107 | |
| 108 | pre_logits = mixes[:, :hc_mult] * hc_scale[0] + hc_base[:hc_mult] |
| 109 | pre_mix = torch.sigmoid(pre_logits) + hc_pre_eps |
| 110 | |
| 111 | post_logits = ( |
| 112 | mixes[:, hc_mult : 2 * hc_mult] * hc_scale[1] + hc_base[hc_mult : 2 * hc_mult] |
| 113 | ) |
| 114 | post_mix = torch.sigmoid(post_logits) * hc_post_mult_value |
| 115 | |
| 116 | comb_logits = mixes[:, 2 * hc_mult :].view(num_tokens, hc_mult, hc_mult) * hc_scale[ |
| 117 | 2 |
| 118 | ] + hc_base[2 * hc_mult :].view(1, hc_mult, hc_mult) |
| 119 | comb_mix = torch.softmax(comb_logits, dim=-1) + hc_sinkhorn_eps |
| 120 | comb_mix = comb_mix / (comb_mix.sum(dim=-2, keepdim=True) + hc_sinkhorn_eps) |
| 121 | for _ in range(sinkhorn_repeat - 1): |
| 122 | comb_mix = comb_mix / (comb_mix.sum(dim=-1, keepdim=True) + hc_sinkhorn_eps) |
| 123 | comb_mix = comb_mix / (comb_mix.sum(dim=-2, keepdim=True) + hc_sinkhorn_eps) |
| 124 | |
| 125 | layer_input = torch.sum( |
| 126 | pre_mix.unsqueeze(-1) * residual_flat.to(torch.float32), dim=1 |
| 127 | ).to(torch.bfloat16) |
| 128 | return ( |
| 129 | post_mix.view(*outer_shape, hc_mult, 1), |
| 130 | comb_mix.view(*outer_shape, hc_mult, hc_mult), |
| 131 | layer_input.view(*outer_shape, hidden_size), |
| 132 | ) |
| 133 | |
| 134 | |
| 135 | def mhc_post_torch( |
| 136 | x: torch.Tensor, |
| 137 | residual: torch.Tensor, |
| 138 | post_layer_mix: torch.Tensor, |
| 139 | comb_res_mix: torch.Tensor, |
| 140 | ) -> torch.Tensor: |
| 141 | mixed_residual = torch.einsum( |
| 142 | "...ij,...ih->...jh", |
| 143 | comb_res_mix.to(torch.float32), |
| 144 | residual.to(torch.float32), |
| 145 | ) |
| 146 | post_term = post_layer_mix.to(torch.float32) * x.unsqueeze(-2).to(torch.float32) |
| 147 | return (mixed_residual + post_term).to(residual.dtype) |