| |
| |
| |
| |
| |
| """Explicit Torch reference backend, not the upstream fused serving kernels.""" |
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
|
|
| def prepare_lens_from_mask(mask): |
| return mask.sum(-1, dtype=torch.int32) |
|
|
|
|
| def prepare_cu_seqlens_from_mask(mask): |
| return F.pad(prepare_lens_from_mask(mask).cumsum(0, dtype=torch.int32), (1, 0)) |
|
|
|
|
| def tensor_cache(function): |
| |
| return function |
|
|
|
|
| def rearrange(x, pattern, **axes): |
| if pattern == 'b s ... -> (b s) ...': |
| return x.flatten(0, 1) |
| if pattern == '... (h d) -> ... h d': |
| return x.unflatten(-1, (x.shape[-1] // axes['d'], axes['d'])) |
| if pattern == 'b t h d -> b t (h d)': |
| return x.flatten(-2) |
| raise ValueError('Unsupported reference reshape: ' + pattern) |
|
|
|
|
| def _segments(x, cu_seqlens): |
| if cu_seqlens is None: |
| return [(index, 0, x.shape[1]) for index in range(x.shape[0])] |
| if x.shape[0] != 1: |
| raise ValueError('Packed input must have batch size one') |
| bounds = cu_seqlens.tolist() |
| if bounds[0] != 0 or bounds[-1] != x.shape[1] or any(a > b for a, b in zip(bounds, bounds[1:])): |
| raise ValueError('Invalid packed cumulative lengths') |
| return [(0, a, b) for a, b in zip(bounds, bounds[1:])] |
|
|
|
|
| class ShortConvolution(nn.Conv1d): |
| def __init__(self, hidden_size, kernel_size, activation='silu', bias=False): |
| super().__init__(hidden_size, hidden_size, kernel_size, groups=hidden_size, bias=bias) |
| if activation not in ('silu', 'swish', None): |
| raise ValueError(activation) |
| self.activation = activation |
|
|
| def forward(self, x, cache=None, output_final_state=False, cu_seqlens=None): |
| width = self.kernel_size[0] |
| output = torch.empty_like(x) |
| final = [] |
| for index, (batch, start, end) in enumerate(_segments(x, cu_seqlens)): |
| state = x.new_zeros(self.in_channels, width) if cache is None else cache[index].clone() |
| for offset in range(start, end): |
| state = torch.cat((state[:, 1:], x[batch, offset, :, None]), -1) |
| value = (state.float() * self.weight[:, 0].float()).sum(-1) |
| if self.bias is not None: |
| value = value + self.bias.float() |
| if self.activation is not None: |
| value = F.silu(value) |
| output[batch, offset] = value.to(x.dtype) |
| final.append(state) |
| final = torch.stack(final) |
| if cache is not None: |
| cache.copy_(final) |
| return output, final if output_final_state else None |
|
|
|
|
| class FusedRMSNormGated(nn.Module): |
| def __init__(self, hidden_size, eps=1e-5, activation='sigmoid'): |
| super().__init__() |
| if activation != 'sigmoid': |
| raise ValueError('This Kimi reference requires sigmoid output gating') |
| self.weight = nn.Parameter(torch.ones(hidden_size)) |
| self.eps = eps |
|
|
| def forward(self, x, gate): |
| normalized = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + self.eps) |
| return (normalized * self.weight.float() * gate.float().sigmoid()).to(x.dtype) |
|
|
|
|
| def _reference_kda(q, k, v, g, beta, A_log=None, dt_bias=None, initial_state=None, |
| output_final_state=False, use_qk_l2norm_in_kernel=False, |
| use_gate_in_kernel=False, use_beta_sigmoid_in_kernel=False, |
| lower_bound=None, transpose_state_layout=False, cu_seqlens=None, |
| safe_gate=False): |
| dtype = v.dtype |
| if q.shape[2] != v.shape[2]: |
| raise ValueError('Kimi K3 reference requires matching query/value heads') |
| q, k, v, g, beta = [value.float() for value in (q, k, v, g, beta)] |
| if use_qk_l2norm_in_kernel: |
| |
| q = q * torch.rsqrt(q.square().sum(-1, keepdim=True) + 1e-6) |
| k = k * torch.rsqrt(k.square().sum(-1, keepdim=True) + 1e-6) |
| if use_gate_in_kernel: |
| if dt_bias is not None: |
| g = g + dt_bias.float().view(g.shape[-2:]) |
| if lower_bound is None: |
| g = -A_log.float().exp().view(-1, 1) * F.softplus(g) |
| else: |
| g = lower_bound * (A_log.float().exp().view(-1, 1) * g).sigmoid() |
| if use_beta_sigmoid_in_kernel: |
| beta = beta.sigmoid() |
| q = q * q.shape[-1] ** -0.5 |
| output = torch.empty_like(v) |
| states = [] |
| for index, (batch, start, end) in enumerate(_segments(q, cu_seqlens)): |
| state = q.new_zeros(q.shape[2], q.shape[3], v.shape[3]) |
| if initial_state is not None: |
| state = initial_state[index].float().clone() |
| if transpose_state_layout: |
| state = state.transpose(-1, -2) |
| for offset in range(start, end): |
| key = k[batch, offset] |
| state = state * g[batch, offset].exp().unsqueeze(-1) |
| delta = v[batch, offset] - (key.unsqueeze(-1) * state).sum(-2) |
| state = state + (beta[batch, offset, :, None] * key).unsqueeze(-1) * delta.unsqueeze(-2) |
| output[batch, offset] = (q[batch, offset].unsqueeze(-1) * state).sum(-2) |
| states.append(state.transpose(-1, -2).contiguous() if transpose_state_layout else state) |
| return output.to(dtype), torch.stack(states) if output_final_state else None |
|
|
|
|
| def chunk_kda(**kwargs): |
| """Evaluate chunk-mode API by its exact sequential recurrence on CPU.""" |
| return _reference_kda(**kwargs) |
|
|
|
|
| def fused_recurrent_kda(**kwargs): |
| """Evaluate recurrent-mode API with the same reference state convention.""" |
| return _reference_kda(**kwargs) |
|
|
