wehub-resource-sync
79 lines
2.7 KiB
Python
79 lines
2.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from einops import rearrange, repeat
|
|
|
|
|
|
def selective_state_update_ref(
|
|
state, x, dt, A, B, C, D=None, z=None, dt_bias=None, dt_softplus=False
|
|
):
|
|
"""
|
|
Argument:
|
|
state: (batch, dim, dstate) or (batch, nheads, dim, dstate)
|
|
x: (batch, dim) or (batch, nheads, dim)
|
|
dt: (batch, dim) or (batch, nheads, dim)
|
|
A: (dim, dstate) or (nheads, dim, dstate)
|
|
B: (batch, dstate) or (batch, ngroups, dstate)
|
|
C: (batch, dstate) or (batch, ngroups, dstate)
|
|
D: (dim,) or (nheads, dim)
|
|
z: (batch, dim) or (batch, nheads, dim)
|
|
dt_bias: (dim,) or (nheads, dim)
|
|
Return:
|
|
out: (batch, dim) or (batch, nheads, dim)
|
|
"""
|
|
has_heads = state.dim() > 3
|
|
if state.dim() == 3:
|
|
state = state.unsqueeze(1)
|
|
if x.dim() == 2:
|
|
x = x.unsqueeze(1)
|
|
if dt.dim() == 2:
|
|
dt = dt.unsqueeze(1)
|
|
if A.dim() == 2:
|
|
A = A.unsqueeze(0)
|
|
if B.dim() == 2:
|
|
B = B.unsqueeze(1)
|
|
if C.dim() == 2:
|
|
C = C.unsqueeze(1)
|
|
if D is not None and D.dim() == 1:
|
|
D = D.unsqueeze(0)
|
|
if z is not None and z.dim() == 2:
|
|
z = z.unsqueeze(1)
|
|
if dt_bias is not None and dt_bias.dim() == 1:
|
|
dt_bias = dt_bias.unsqueeze(0)
|
|
batch, nheads, dim, dstate = state.shape
|
|
assert x.shape == (batch, nheads, dim)
|
|
assert dt.shape == x.shape
|
|
assert A.shape == (nheads, dim, dstate)
|
|
ngroups = B.shape[1]
|
|
assert nheads % ngroups == 0, "nheads must be divisible by ngroups"
|
|
assert B.shape == (batch, ngroups, dstate)
|
|
assert C.shape == B.shape
|
|
if D is not None:
|
|
assert D.shape == (nheads, dim)
|
|
if z is not None:
|
|
assert z.shape == x.shape
|
|
if dt_bias is not None:
|
|
assert dt_bias.shape == (nheads, dim)
|
|
dt = dt + dt_bias
|
|
dt = F.softplus(dt) if dt_softplus else dt
|
|
dA = torch.exp(
|
|
rearrange(dt, "b h d -> b h d 1") * A
|
|
) # (batch, nheads, dim, dstate)
|
|
B = repeat(B, "b g n -> b (g h) n", h=nheads // ngroups) # (batch, nheads, dstate)
|
|
C = repeat(C, "b g n -> b (g h) n", h=nheads // ngroups) # (batch, nheads, dstate)
|
|
dB = rearrange(dt, "b h d -> b h d 1") * rearrange(
|
|
B, "b h n -> b h 1 n"
|
|
) # (batch, nheads, dim, dstate)
|
|
state.copy_(
|
|
state * dA + dB * rearrange(x, "b h d -> b h d 1")
|
|
) # (batch, dim, dstate
|
|
out = torch.einsum("bhdn,bhn->bhd", state.to(C.dtype), C)
|
|
if D is not None:
|
|
out += (x * D).to(out.dtype)
|
|
out = (out if z is None else out * F.silu(z)).to(x.dtype)
|
|
if not has_heads:
|
|
out = out.squeeze(1)
|
|
return out
|